Merge pull request #63 from unslothai/feature/chat-openai-integration
feat(chat): integrate backend chat runtime + model load flow
This commit is contained in:
Wasim Yousef Said
GitHub
commit
ddfa6c40c6
16 changed files with 1186 additions and 329 deletions
|
|
@ -10,7 +10,7 @@ export function AppProvider({ children }: AppProviderProps) {
|
|||
return (
|
||||
<ThemeProvider attribute="class" defaultTheme="light">
|
||||
{children}
|
||||
<Toaster />
|
||||
<Toaster position="top-right" />
|
||||
</ThemeProvider>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,14 +1,22 @@
|
|||
"use client";
|
||||
|
||||
import { Input } from "@/components/ui/input";
|
||||
import {
|
||||
Popover,
|
||||
PopoverContent,
|
||||
PopoverTrigger,
|
||||
} from "@/components/ui/popover";
|
||||
import { cn } from "@/lib/utils";
|
||||
import { ArrowDown01Icon, Logout01Icon } from "@hugeicons/core-free-icons";
|
||||
import { Spinner } from "@/components/ui/spinner";
|
||||
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
|
||||
import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks";
|
||||
import { cn, formatCompact } from "@/lib/utils";
|
||||
import {
|
||||
ArrowDown01Icon,
|
||||
Logout01Icon,
|
||||
Search01Icon,
|
||||
} from "@hugeicons/core-free-icons";
|
||||
import { HugeiconsIcon } from "@hugeicons/react";
|
||||
import { type ReactNode, useState } from "react";
|
||||
import { type ReactNode, useMemo, useState } from "react";
|
||||
|
||||
export interface ModelOption {
|
||||
id: string;
|
||||
|
|
@ -17,11 +25,22 @@ export interface ModelOption {
|
|||
icon?: ReactNode;
|
||||
}
|
||||
|
||||
export interface LoraModelOption extends ModelOption {
|
||||
baseModel?: string;
|
||||
updatedAt?: number;
|
||||
}
|
||||
|
||||
export interface ModelSelectorChangeMeta {
|
||||
source: "hub" | "lora";
|
||||
isLora: boolean;
|
||||
}
|
||||
|
||||
interface ModelSelectorProps {
|
||||
models: ModelOption[];
|
||||
loraModels?: LoraModelOption[];
|
||||
value?: string;
|
||||
defaultValue?: string;
|
||||
onValueChange?: (value: string) => void;
|
||||
onValueChange?: (value: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
variant?: "outline" | "ghost" | "muted";
|
||||
size?: "sm" | "default" | "lg";
|
||||
|
|
@ -29,7 +48,9 @@ interface ModelSelectorProps {
|
|||
contentClassName?: string;
|
||||
}
|
||||
|
||||
// --- Composable sub-components ---
|
||||
function dedupe(values: string[]): string[] {
|
||||
return [...new Set(values.filter(Boolean))];
|
||||
}
|
||||
|
||||
function ModelSelectorTrigger({
|
||||
currentModel,
|
||||
|
|
@ -63,15 +84,11 @@ function ModelSelectorTrigger({
|
|||
{isLoaded && (
|
||||
<span className="size-2 shrink-0 rounded-full bg-emerald-500" />
|
||||
)}
|
||||
<span
|
||||
className={isLoaded ? "text-foreground" : "text-muted-foreground"}
|
||||
>
|
||||
{currentModel?.name ?? "Select a model\u2026"}
|
||||
<span className={isLoaded ? "text-foreground" : "text-muted-foreground"}>
|
||||
{currentModel?.name ?? "Select model..."}
|
||||
</span>
|
||||
{currentModel?.description && (
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{currentModel.description}
|
||||
</span>
|
||||
<span className="text-muted-foreground text-xs">{currentModel.description}</span>
|
||||
)}
|
||||
<HugeiconsIcon
|
||||
icon={ArrowDown01Icon}
|
||||
|
|
@ -82,95 +99,329 @@ function ModelSelectorTrigger({
|
|||
);
|
||||
}
|
||||
|
||||
function ListLabel({ children }: { children: ReactNode }) {
|
||||
return (
|
||||
<div className="px-2.5 py-1.5 text-[10px] font-semibold uppercase tracking-wider text-muted-foreground">
|
||||
{children}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ModelRow({
|
||||
label,
|
||||
meta,
|
||||
selected,
|
||||
onClick,
|
||||
}: {
|
||||
label: string;
|
||||
meta?: string;
|
||||
selected?: boolean;
|
||||
onClick: () => void;
|
||||
}) {
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
onClick={onClick}
|
||||
className={cn(
|
||||
"flex w-full items-center justify-between gap-2 rounded-md px-2.5 py-1.5 text-left text-sm transition-colors hover:bg-accent",
|
||||
selected && "bg-accent/60",
|
||||
)}
|
||||
>
|
||||
<span className="min-w-0 flex-1 truncate">{label}</span>
|
||||
{meta ? (
|
||||
<span className="shrink-0 text-[10px] text-muted-foreground">{meta}</span>
|
||||
) : null}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
function HubModelPicker({
|
||||
models,
|
||||
value,
|
||||
onSelect,
|
||||
}: {
|
||||
models: ModelOption[];
|
||||
value?: string;
|
||||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
}) {
|
||||
const [query, setQuery] = useState("");
|
||||
const debouncedQuery = useDebouncedValue(query);
|
||||
const { results, isLoading, isLoadingMore, fetchMore } = useHfModelSearch(
|
||||
debouncedQuery,
|
||||
);
|
||||
|
||||
const recommendedIds = useMemo(
|
||||
() => dedupe([...models.map((model) => model.id), value ?? ""]),
|
||||
[models, value],
|
||||
);
|
||||
|
||||
const showHfSection = debouncedQuery.trim().length > 0;
|
||||
const recommendedSet = useMemo(
|
||||
() => new Set(recommendedIds),
|
||||
[recommendedIds],
|
||||
);
|
||||
|
||||
const hfIds = useMemo(() => {
|
||||
if (!showHfSection) {
|
||||
return [];
|
||||
}
|
||||
return results
|
||||
.map((result) => result.id)
|
||||
.filter((id) => !recommendedSet.has(id));
|
||||
}, [recommendedSet, results, showHfSection]);
|
||||
|
||||
const metricsById = useMemo(
|
||||
() =>
|
||||
new Map(
|
||||
results.map((result) => [
|
||||
result.id,
|
||||
result.totalParams
|
||||
? formatCompact(result.totalParams)
|
||||
: `↓${formatCompact(result.downloads)}`,
|
||||
]),
|
||||
),
|
||||
[results],
|
||||
);
|
||||
|
||||
const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore, results.length);
|
||||
|
||||
return (
|
||||
<div className="space-y-2">
|
||||
<div className="relative">
|
||||
<HugeiconsIcon
|
||||
icon={Search01Icon}
|
||||
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
|
||||
/>
|
||||
<Input
|
||||
value={query}
|
||||
onChange={(event) => setQuery(event.target.value)}
|
||||
placeholder="Search Hugging Face models"
|
||||
className="h-9 pl-8 pr-8"
|
||||
/>
|
||||
{isLoading && (
|
||||
<Spinner className="pointer-events-none absolute right-2.5 top-2.5 size-4 text-muted-foreground" />
|
||||
)}
|
||||
</div>
|
||||
|
||||
<div
|
||||
ref={scrollRef}
|
||||
className="max-h-64 overflow-y-auto"
|
||||
>
|
||||
<div className="p-1">
|
||||
{!showHfSection ? (
|
||||
<>
|
||||
<ListLabel>Recommended</ListLabel>
|
||||
{recommendedIds.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No default models.
|
||||
</div>
|
||||
) : (
|
||||
recommendedIds.map((id) => (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
selected={value === id}
|
||||
onClick={() => onSelect(id, { source: "hub", isLora: false })}
|
||||
/>
|
||||
))
|
||||
)}
|
||||
</>
|
||||
) : null}
|
||||
|
||||
{showHfSection ? (
|
||||
<>
|
||||
<ListLabel>Hugging Face</ListLabel>
|
||||
{hfIds.length === 0 && !isLoading ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">
|
||||
No matching models.
|
||||
</div>
|
||||
) : (
|
||||
hfIds.map((id) => (
|
||||
<ModelRow
|
||||
key={id}
|
||||
label={id}
|
||||
meta={metricsById.get(id)}
|
||||
selected={value === id}
|
||||
onClick={() => onSelect(id, { source: "hub", isLora: false })}
|
||||
/>
|
||||
))
|
||||
)}
|
||||
<div ref={sentinelRef} className="h-px" />
|
||||
{isLoadingMore ? (
|
||||
<div className="flex items-center justify-center py-2">
|
||||
<Spinner className="size-3.5 text-muted-foreground" />
|
||||
</div>
|
||||
) : null}
|
||||
</>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function LoraModelPicker({
|
||||
loraModels,
|
||||
value,
|
||||
onSelect,
|
||||
}: {
|
||||
loraModels: LoraModelOption[];
|
||||
value?: string;
|
||||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
}) {
|
||||
const [query, setQuery] = useState("");
|
||||
|
||||
const normalized = useMemo(
|
||||
() =>
|
||||
loraModels
|
||||
.map((model) => ({
|
||||
...model,
|
||||
baseModel: model.baseModel || model.description || "Unknown base model",
|
||||
}))
|
||||
.sort((a, b) => {
|
||||
const aTime = a.updatedAt ?? -1;
|
||||
const bTime = b.updatedAt ?? -1;
|
||||
if (aTime !== bTime) {
|
||||
return bTime - aTime;
|
||||
}
|
||||
const baseCmp = a.baseModel.localeCompare(b.baseModel);
|
||||
if (baseCmp !== 0) {
|
||||
return baseCmp;
|
||||
}
|
||||
return a.name.localeCompare(b.name);
|
||||
}),
|
||||
[loraModels],
|
||||
);
|
||||
|
||||
const grouped = useMemo(() => {
|
||||
const needle = query.trim().toLowerCase();
|
||||
const out = new Map<string, LoraModelOption[]>();
|
||||
|
||||
for (const model of normalized) {
|
||||
const searchText = `${model.name} ${model.baseModel} ${model.id}`.toLowerCase();
|
||||
if (needle && !searchText.includes(needle)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const key = model.baseModel || "Unknown base model";
|
||||
const prev = out.get(key) ?? [];
|
||||
prev.push(model);
|
||||
out.set(key, prev);
|
||||
}
|
||||
|
||||
return [...out.entries()].sort((a, b) => {
|
||||
const aLatest = Math.max(...a[1].map((model) => model.updatedAt ?? -1));
|
||||
const bLatest = Math.max(...b[1].map((model) => model.updatedAt ?? -1));
|
||||
if (aLatest !== bLatest) {
|
||||
return bLatest - aLatest;
|
||||
}
|
||||
return a[0].localeCompare(b[0]);
|
||||
});
|
||||
}, [normalized, query]);
|
||||
|
||||
return (
|
||||
<div className="space-y-2">
|
||||
<div className="relative">
|
||||
<HugeiconsIcon
|
||||
icon={Search01Icon}
|
||||
className="pointer-events-none absolute left-2.5 top-2.5 size-4 text-muted-foreground"
|
||||
/>
|
||||
<Input
|
||||
value={query}
|
||||
onChange={(event) => setQuery(event.target.value)}
|
||||
placeholder="Search local adapters"
|
||||
className="h-9 pl-8"
|
||||
/>
|
||||
</div>
|
||||
|
||||
<div className="max-h-64 overflow-y-auto">
|
||||
<div className="p-1">
|
||||
{grouped.length === 0 ? (
|
||||
<div className="px-2.5 py-2 text-xs text-muted-foreground">No adapters found.</div>
|
||||
) : (
|
||||
grouped.map(([baseModel, adapters], index) => (
|
||||
<div key={baseModel}>
|
||||
{index > 0 ? <div className="my-1" /> : null}
|
||||
<ListLabel>{baseModel}</ListLabel>
|
||||
{adapters.map((adapter) => (
|
||||
<ModelRow
|
||||
key={adapter.id}
|
||||
label={adapter.name}
|
||||
meta="LoRA"
|
||||
selected={value === adapter.id}
|
||||
onClick={() => onSelect(adapter.id, { source: "lora", isLora: true })}
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
))
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ModelSelectorContent({
|
||||
models,
|
||||
loraModels,
|
||||
value,
|
||||
onSelect,
|
||||
onEject,
|
||||
className,
|
||||
}: {
|
||||
models: ModelOption[];
|
||||
loraModels: LoraModelOption[];
|
||||
value?: string;
|
||||
onSelect: (id: string) => void;
|
||||
onSelect: (id: string, meta: ModelSelectorChangeMeta) => void;
|
||||
onEject?: () => void;
|
||||
className?: string;
|
||||
}) {
|
||||
const hasSelection = Boolean(value);
|
||||
|
||||
return (
|
||||
<PopoverContent
|
||||
align="start"
|
||||
className={cn("w-auto min-w-[280px] gap-0 p-1", className)}
|
||||
className={cn("w-[440px] min-w-[440px] gap-0 p-2", className)}
|
||||
>
|
||||
{models.map((model) => (
|
||||
<ModelSelectorItem
|
||||
key={model.id}
|
||||
model={model}
|
||||
isActive={value === model.id}
|
||||
onSelect={onSelect}
|
||||
onEject={onEject}
|
||||
/>
|
||||
))}
|
||||
<Tabs defaultValue="hub" className="w-full">
|
||||
<TabsList className="mb-2 w-full">
|
||||
<TabsTrigger value="hub">Hub models</TabsTrigger>
|
||||
<TabsTrigger value="lora">Fine-tuned</TabsTrigger>
|
||||
</TabsList>
|
||||
|
||||
<TabsContent value="hub" className="m-0">
|
||||
<HubModelPicker models={models} value={value} onSelect={onSelect} />
|
||||
</TabsContent>
|
||||
|
||||
<TabsContent value="lora" className="m-0">
|
||||
<LoraModelPicker
|
||||
loraModels={loraModels}
|
||||
value={value}
|
||||
onSelect={onSelect}
|
||||
/>
|
||||
</TabsContent>
|
||||
</Tabs>
|
||||
|
||||
{hasSelection && onEject ? (
|
||||
<div className="mt-2 border-t border-border/70 pt-2">
|
||||
<button
|
||||
type="button"
|
||||
onClick={onEject}
|
||||
className="flex w-full items-center justify-center gap-1.5 rounded-md px-2 py-1.5 text-xs text-muted-foreground transition-colors hover:bg-accent hover:text-foreground"
|
||||
title="Eject model"
|
||||
>
|
||||
<HugeiconsIcon icon={Logout01Icon} className="size-3.5" />
|
||||
Eject loaded model
|
||||
</button>
|
||||
</div>
|
||||
) : null}
|
||||
</PopoverContent>
|
||||
);
|
||||
}
|
||||
|
||||
function ModelSelectorItem({
|
||||
model,
|
||||
isActive,
|
||||
onSelect,
|
||||
onEject,
|
||||
}: {
|
||||
model: ModelOption;
|
||||
isActive: boolean;
|
||||
onSelect: (id: string) => void;
|
||||
onEject?: () => void;
|
||||
}) {
|
||||
return (
|
||||
<button
|
||||
type="button"
|
||||
aria-pressed={isActive}
|
||||
onClick={() => onSelect(model.id)}
|
||||
className={cn(
|
||||
"group flex w-full cursor-pointer items-center gap-2.5 rounded-md px-3 py-2 text-left text-sm transition-colors hover:bg-accent",
|
||||
isActive && "bg-accent/50",
|
||||
)}
|
||||
>
|
||||
<span
|
||||
className={cn(
|
||||
"size-2 shrink-0 rounded-full",
|
||||
isActive ? "bg-emerald-500" : "bg-transparent",
|
||||
)}
|
||||
/>
|
||||
{model.icon && <span className="shrink-0">{model.icon}</span>}
|
||||
<div className="min-w-0 flex-1">
|
||||
<div className="truncate text-sm">{model.name}</div>
|
||||
{model.description && (
|
||||
<div className="truncate text-xs text-muted-foreground">
|
||||
{model.description}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
{isActive && onEject && (
|
||||
<button
|
||||
type="button"
|
||||
onClick={(e) => {
|
||||
e.stopPropagation();
|
||||
onEject();
|
||||
}}
|
||||
className="flex items-center gap-1 rounded px-1.5 py-0.5 text-[10px] text-muted-foreground opacity-0 transition-opacity group-hover:opacity-100 hover:bg-muted hover:text-foreground"
|
||||
title="Eject model"
|
||||
>
|
||||
<HugeiconsIcon icon={Logout01Icon} className="size-3" />
|
||||
Eject
|
||||
</button>
|
||||
)}
|
||||
</button>
|
||||
);
|
||||
}
|
||||
|
||||
// --- Main component ---
|
||||
|
||||
export function ModelSelector({
|
||||
models,
|
||||
loraModels = [],
|
||||
value,
|
||||
defaultValue,
|
||||
onValueChange,
|
||||
|
|
@ -182,13 +433,31 @@ export function ModelSelector({
|
|||
}: ModelSelectorProps) {
|
||||
const [open, setOpen] = useState(false);
|
||||
const [uncontrolled, setUncontrolled] = useState(defaultValue ?? "");
|
||||
|
||||
const selected = value ?? uncontrolled;
|
||||
const isLoaded = selected !== "";
|
||||
const currentModel = models.find((m) => m.id === selected);
|
||||
|
||||
function handleSelect(id: string) {
|
||||
const optionById = useMemo(() => {
|
||||
const all = new Map<string, ModelOption>();
|
||||
for (const model of models) {
|
||||
all.set(model.id, model);
|
||||
}
|
||||
for (const lora of loraModels) {
|
||||
all.set(lora.id, {
|
||||
...lora,
|
||||
description: lora.baseModel || lora.description,
|
||||
});
|
||||
}
|
||||
return all;
|
||||
}, [loraModels, models]);
|
||||
|
||||
const currentModel = selected
|
||||
? optionById.get(selected) ?? { id: selected, name: selected }
|
||||
: undefined;
|
||||
|
||||
function handleSelect(id: string, meta: ModelSelectorChangeMeta) {
|
||||
if (onValueChange) {
|
||||
onValueChange(id);
|
||||
onValueChange(id, meta);
|
||||
} else {
|
||||
setUncontrolled(id);
|
||||
}
|
||||
|
|
@ -211,6 +480,7 @@ export function ModelSelector({
|
|||
/>
|
||||
<ModelSelectorContent
|
||||
models={models}
|
||||
loraModels={loraModels}
|
||||
value={selected}
|
||||
onSelect={handleSelect}
|
||||
onEject={onEject ? handleEject : undefined}
|
||||
|
|
@ -220,7 +490,5 @@ export function ModelSelector({
|
|||
);
|
||||
}
|
||||
|
||||
// Composable exports
|
||||
ModelSelector.Trigger = ModelSelectorTrigger;
|
||||
ModelSelector.Content = ModelSelectorContent;
|
||||
ModelSelector.Item = ModelSelectorItem;
|
||||
|
|
|
|||
|
|
@ -7,7 +7,9 @@ import { MarkdownText } from "@/components/assistant-ui/markdown-text";
|
|||
import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
|
||||
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
|
||||
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
|
||||
import { AnimatedShinyText } from "@/components/ui/animated-shiny-text";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
|
||||
import { cn } from "@/lib/utils";
|
||||
import {
|
||||
ActionBarMorePrimitive,
|
||||
|
|
@ -20,6 +22,8 @@ import {
|
|||
SuggestionPrimitive,
|
||||
ThreadPrimitive,
|
||||
useAui,
|
||||
useAuiEvent,
|
||||
useAuiState,
|
||||
} from "@assistant-ui/react";
|
||||
import { motion } from "framer-motion";
|
||||
import {
|
||||
|
|
@ -36,7 +40,7 @@ import {
|
|||
RefreshCwIcon,
|
||||
SquareIcon,
|
||||
} from "lucide-react";
|
||||
import type { FC } from "react";
|
||||
import { type FC, useRef } from "react";
|
||||
|
||||
export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
|
||||
hideComposer,
|
||||
|
|
@ -69,6 +73,7 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
|
|||
|
||||
<ThreadPrimitive.ViewportFooter className="aui-thread-viewport-footer sticky bottom-0 mt-auto flex w-full flex-col gap-4 overflow-visible bg-background pb-4 md:pb-4 before:pointer-events-none before:absolute before:inset-x-0 before:bottom-full before:h-20 before:bg-gradient-to-t before:from-background before:to-transparent">
|
||||
<ThreadScrollToBottom />
|
||||
<WarmupIndicator />
|
||||
<AuiIf condition={({ thread }) => !thread.isEmpty}>
|
||||
{!hideComposer && <ComposerAnimated />}
|
||||
</AuiIf>
|
||||
|
|
@ -78,6 +83,28 @@ export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
|
|||
);
|
||||
};
|
||||
|
||||
const WarmupIndicator: FC = () => {
|
||||
const threadId = useAuiState(({ threads }) => threads.mainThreadId);
|
||||
const isRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
const isWarmingUp = useChatRuntimeStore((state) =>
|
||||
Boolean(state.warmingByThreadId[threadId ?? "__default"]),
|
||||
);
|
||||
|
||||
if (!isRunning || !isWarmingUp) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mx-auto -mb-2 w-full max-w-(--thread-max-width) px-2">
|
||||
<div className="inline-flex items-center rounded-full border border-border/60 bg-background/90 px-3 py-1.5 text-xs text-muted-foreground shadow-sm">
|
||||
<AnimatedShinyText className="text-xs">
|
||||
Warming up model...
|
||||
</AnimatedShinyText>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
const ThreadScrollToBottom: FC = () => {
|
||||
return (
|
||||
<ThreadPrimitive.ScrollToBottom asChild={true}>
|
||||
|
|
@ -93,13 +120,26 @@ const ThreadScrollToBottom: FC = () => {
|
|||
};
|
||||
|
||||
const SuggestionItem: FC = () => {
|
||||
const aui = useAui();
|
||||
const prompt = useAuiState(({ suggestion }) => suggestion.prompt);
|
||||
const isDisabled = useAuiState(({ thread }) => thread.isDisabled);
|
||||
const isRunning = useAuiState(({ thread }) => thread.isRunning);
|
||||
|
||||
return (
|
||||
<SuggestionPrimitive.Trigger
|
||||
send={true}
|
||||
<button
|
||||
type="button"
|
||||
onClick={() => {
|
||||
if (!isDisabled && !isRunning) {
|
||||
aui.thread().append(prompt);
|
||||
aui.composer().setText("");
|
||||
return;
|
||||
}
|
||||
aui.composer().setText(prompt);
|
||||
}}
|
||||
className="fade-in slide-in-from-bottom-1 animate-in cursor-pointer corner-squircle rounded-xl border bg-background px-4 py-2.5 text-left text-sm text-foreground shadow-sm transition-colors duration-150 hover:bg-accent"
|
||||
>
|
||||
<SuggestionPrimitive.Title />
|
||||
</SuggestionPrimitive.Trigger>
|
||||
</button>
|
||||
);
|
||||
};
|
||||
|
||||
|
|
@ -358,6 +398,15 @@ const UserActionBar: FC = () => {
|
|||
|
||||
const EditComposer: FC = () => {
|
||||
const aui = useAui();
|
||||
const resendAfterCancelRef = useRef(false);
|
||||
|
||||
useAuiEvent("thread.runEnd", () => {
|
||||
if (!resendAfterCancelRef.current) {
|
||||
return;
|
||||
}
|
||||
resendAfterCancelRef.current = false;
|
||||
aui.composer().send();
|
||||
});
|
||||
|
||||
return (
|
||||
<MessagePrimitive.Root className="aui-edit-composer-wrapper mx-auto flex w-full max-w-(--thread-max-width) flex-col px-2 py-3">
|
||||
|
|
@ -384,7 +433,9 @@ const EditComposer: FC = () => {
|
|||
}
|
||||
|
||||
if (aui.thread().getState().isRunning) {
|
||||
resendAfterCancelRef.current = true;
|
||||
aui.thread().cancelRun();
|
||||
return;
|
||||
}
|
||||
aui.composer().send();
|
||||
}}
|
||||
|
|
|
|||
|
|
@ -1,121 +0,0 @@
|
|||
import type { ChatModelAdapter, ChatModelRunResult } from "@assistant-ui/react";
|
||||
|
||||
const API = import.meta.env.VITE_INFERENCE_URL || "/api/chat/generate";
|
||||
type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
|
||||
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
|
||||
type RunMessage = RunMessages[number];
|
||||
|
||||
function collectTextParts(message: RunMessage): string[] {
|
||||
const textParts = message.content
|
||||
.filter((c) => c.type === "text")
|
||||
.map((c) => c.text);
|
||||
|
||||
if ("attachments" in message && (message.attachments?.length ?? 0) > 0) {
|
||||
for (const att of message.attachments ?? []) {
|
||||
for (const part of att.content ?? []) {
|
||||
if (part.type === "text") {
|
||||
textParts.push(part.text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return textParts;
|
||||
}
|
||||
|
||||
function messageToPayload(message: RunMessage): {
|
||||
role: string;
|
||||
content: string;
|
||||
} {
|
||||
return {
|
||||
role: message.role,
|
||||
content: collectTextParts(message).join("\n"),
|
||||
};
|
||||
}
|
||||
|
||||
function makeBody(messages: RunMessages): string {
|
||||
const payloadMessages: Array<{ role: string; content: string }> = [];
|
||||
for (const message of messages) {
|
||||
payloadMessages.push(messageToPayload(message));
|
||||
}
|
||||
return JSON.stringify({ messages: payloadMessages });
|
||||
}
|
||||
|
||||
export function parseThinkTags(raw: string): ChatModelRunResult["content"] {
|
||||
const parts: ContentPart[] = [];
|
||||
const thinkStart = raw.indexOf("<think>");
|
||||
if (thinkStart === -1) {
|
||||
if (raw) {
|
||||
parts.push({ type: "text", text: raw });
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
const before = raw.slice(0, thinkStart);
|
||||
if (before.trim()) {
|
||||
parts.push({ type: "text", text: before });
|
||||
}
|
||||
|
||||
const thinkEnd = raw.indexOf("</think>");
|
||||
if (thinkEnd === -1) {
|
||||
const reasoning = raw.slice(thinkStart + 7);
|
||||
if (reasoning) {
|
||||
parts.push({ type: "reasoning", text: reasoning });
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
const reasoning = raw.slice(thinkStart + 7, thinkEnd);
|
||||
if (reasoning) {
|
||||
parts.push({ type: "reasoning", text: reasoning });
|
||||
}
|
||||
|
||||
const after = raw.slice(thinkEnd + 8);
|
||||
if (after) {
|
||||
parts.push({ type: "text", text: after });
|
||||
}
|
||||
return parts;
|
||||
}
|
||||
|
||||
export function createStreamAdapter(apiUrl: string = API): ChatModelAdapter {
|
||||
return {
|
||||
// biome-ignore lint/complexity/noExcessiveCognitiveComplexity: stream loop ok
|
||||
async *run({ messages, abortSignal }) {
|
||||
const res = await fetch(apiUrl, {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: makeBody(messages),
|
||||
signal: abortSignal,
|
||||
});
|
||||
const reader = res.body?.getReader();
|
||||
if (!reader) {
|
||||
throw new Error("Response body is empty");
|
||||
}
|
||||
const decoder = new TextDecoder();
|
||||
let text = "";
|
||||
let reasoningStart: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
break;
|
||||
}
|
||||
text += decoder.decode(value, { stream: true });
|
||||
const parts = parseThinkTags(text) ?? [];
|
||||
|
||||
if (parts.some((p) => p.type === "reasoning") && !reasoningStart) {
|
||||
reasoningStart = Date.now();
|
||||
}
|
||||
if (text.includes("</think>") && reasoningStart && !reasoningDuration) {
|
||||
reasoningDuration = Math.round((Date.now() - reasoningStart) / 1000);
|
||||
}
|
||||
|
||||
if (parts.length > 0) {
|
||||
yield {
|
||||
content: parts,
|
||||
metadata: { custom: { reasoningDuration } },
|
||||
};
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
127
studio/frontend/src/features/chat/api/chat-adapter.ts
Normal file
127
studio/frontend/src/features/chat/api/chat-adapter.ts
Normal file
|
|
@ -0,0 +1,127 @@
|
|||
import type { ChatModelAdapter } from "@assistant-ui/react";
|
||||
import { streamChatCompletions } from "./chat-api";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import {
|
||||
hasClosedThinkTag,
|
||||
parseAssistantContent,
|
||||
} from "../utils/parse-assistant-content";
|
||||
|
||||
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
|
||||
type RunMessage = RunMessages[number];
|
||||
|
||||
function collectTextParts(message: RunMessage): string[] {
|
||||
const textParts = message.content
|
||||
.filter((part) => part.type === "text")
|
||||
.map((part) => part.text);
|
||||
|
||||
if ("attachments" in message && (message.attachments?.length ?? 0) > 0) {
|
||||
for (const attachment of message.attachments ?? []) {
|
||||
for (const part of attachment.content ?? []) {
|
||||
if (part.type === "text") {
|
||||
textParts.push(part.text);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return textParts;
|
||||
}
|
||||
|
||||
function toOpenAIMessage(message: RunMessage): {
|
||||
role: "system" | "user" | "assistant";
|
||||
content: string;
|
||||
} | null {
|
||||
if (
|
||||
message.role !== "system" &&
|
||||
message.role !== "user" &&
|
||||
message.role !== "assistant"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
role: message.role,
|
||||
content: collectTextParts(message).join("\n"),
|
||||
};
|
||||
}
|
||||
|
||||
export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
||||
return {
|
||||
async *run({ messages, abortSignal, unstable_threadId }) {
|
||||
const state = useChatRuntimeStore.getState();
|
||||
const { params } = state;
|
||||
|
||||
if (!params.checkpoint) {
|
||||
throw new Error("Load a model first.");
|
||||
}
|
||||
|
||||
const outboundMessages = messages
|
||||
.map(toOpenAIMessage)
|
||||
.filter((message): message is NonNullable<typeof message> =>
|
||||
Boolean(message),
|
||||
);
|
||||
|
||||
if (params.systemPrompt.trim()) {
|
||||
outboundMessages.unshift({
|
||||
role: "system",
|
||||
content: params.systemPrompt.trim(),
|
||||
});
|
||||
}
|
||||
|
||||
const threadKey = unstable_threadId || "__default";
|
||||
let waitingFirstChunk = true;
|
||||
useChatRuntimeStore.getState().setThreadWarming(threadKey, true);
|
||||
let cumulativeText = "";
|
||||
let reasoningStartAt: number | null = null;
|
||||
let reasoningDuration = 0;
|
||||
|
||||
try {
|
||||
const stream = streamChatCompletions(
|
||||
{
|
||||
model: params.checkpoint,
|
||||
messages: outboundMessages,
|
||||
stream: true,
|
||||
temperature: params.temperature,
|
||||
top_p: params.topP,
|
||||
max_tokens: params.maxTokens,
|
||||
top_k: params.topK,
|
||||
repetition_penalty: params.repetitionPenalty,
|
||||
},
|
||||
abortSignal,
|
||||
);
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices?.[0]?.delta?.content;
|
||||
if (!delta) {
|
||||
continue;
|
||||
}
|
||||
if (waitingFirstChunk) {
|
||||
waitingFirstChunk = false;
|
||||
useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
|
||||
}
|
||||
|
||||
cumulativeText += delta;
|
||||
const parts = parseAssistantContent(cumulativeText);
|
||||
|
||||
if (parts.some((part) => part.type === "reasoning") && !reasoningStartAt) {
|
||||
reasoningStartAt = Date.now();
|
||||
}
|
||||
if (hasClosedThinkTag(cumulativeText) && reasoningStartAt && !reasoningDuration) {
|
||||
reasoningDuration = Math.round((Date.now() - reasoningStartAt) / 1000);
|
||||
}
|
||||
|
||||
if (parts.length > 0) {
|
||||
yield {
|
||||
content: parts,
|
||||
metadata: { custom: { reasoningDuration } },
|
||||
};
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
if (waitingFirstChunk) {
|
||||
useChatRuntimeStore.getState().setThreadWarming(threadKey, false);
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
}
|
||||
146
studio/frontend/src/features/chat/api/chat-api.ts
Normal file
146
studio/frontend/src/features/chat/api/chat-api.ts
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
import { authFetch } from "@/features/auth";
|
||||
import type {
|
||||
InferenceStatusResponse,
|
||||
ListLorasResponse,
|
||||
ListModelsResponse,
|
||||
LoadModelRequest,
|
||||
LoadModelResponse,
|
||||
OpenAIChatChunk,
|
||||
OpenAIChatCompletionsRequest,
|
||||
UnloadModelRequest,
|
||||
} from "../types/api";
|
||||
|
||||
function parseErrorText(status: number, body: unknown): string {
|
||||
if (
|
||||
body &&
|
||||
typeof body === "object" &&
|
||||
"detail" in body &&
|
||||
typeof body.detail === "string"
|
||||
) {
|
||||
return body.detail;
|
||||
}
|
||||
if (
|
||||
body &&
|
||||
typeof body === "object" &&
|
||||
"message" in body &&
|
||||
typeof body.message === "string"
|
||||
) {
|
||||
return body.message;
|
||||
}
|
||||
return `Request failed (${status})`;
|
||||
}
|
||||
|
||||
async function parseJsonOrThrow<T>(response: Response): Promise<T> {
|
||||
const body = await response.json().catch(() => null);
|
||||
if (!response.ok) {
|
||||
throw new Error(parseErrorText(response.status, body));
|
||||
}
|
||||
return body as T;
|
||||
}
|
||||
|
||||
export async function listModels(): Promise<ListModelsResponse> {
|
||||
const response = await authFetch("/api/models/list");
|
||||
return parseJsonOrThrow<ListModelsResponse>(response);
|
||||
}
|
||||
|
||||
export async function listLoras(outputsDir = "./outputs"): Promise<ListLorasResponse> {
|
||||
const query = new URLSearchParams({ outputs_dir: outputsDir }).toString();
|
||||
const response = await authFetch(`/api/models/loras?${query}`);
|
||||
return parseJsonOrThrow<ListLorasResponse>(response);
|
||||
}
|
||||
|
||||
export async function getInferenceStatus(): Promise<InferenceStatusResponse> {
|
||||
const response = await authFetch("/api/inference/status");
|
||||
return parseJsonOrThrow<InferenceStatusResponse>(response);
|
||||
}
|
||||
|
||||
export async function loadModel(
|
||||
payload: LoadModelRequest,
|
||||
): Promise<LoadModelResponse> {
|
||||
const response = await authFetch("/api/inference/load", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
return parseJsonOrThrow<LoadModelResponse>(response);
|
||||
}
|
||||
|
||||
export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
|
||||
const response = await authFetch("/api/inference/unload", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
await parseJsonOrThrow<unknown>(response);
|
||||
}
|
||||
|
||||
function parseSseEvent(rawEvent: string): string[] {
|
||||
const dataLines: string[] = [];
|
||||
for (const line of rawEvent.split(/\r?\n/)) {
|
||||
if (line.startsWith("data:")) {
|
||||
dataLines.push(line.slice(5).trimStart());
|
||||
}
|
||||
}
|
||||
return dataLines;
|
||||
}
|
||||
|
||||
export async function* streamChatCompletions(
|
||||
payload: OpenAIChatCompletionsRequest,
|
||||
signal: AbortSignal,
|
||||
): AsyncGenerator<OpenAIChatChunk> {
|
||||
const response = await authFetch("/api/inference/chat/completions", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(payload),
|
||||
signal,
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const body = await response.json().catch(() => null);
|
||||
throw new Error(parseErrorText(response.status, body));
|
||||
}
|
||||
|
||||
if (!response.body) {
|
||||
throw new Error("Stream response missing body");
|
||||
}
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
break;
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true });
|
||||
|
||||
let separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
while (separatorIndex >= 0) {
|
||||
const rawEvent = buffer.slice(0, separatorIndex);
|
||||
const separatorLength = buffer[separatorIndex] === "\r" ? 4 : 2;
|
||||
buffer = buffer.slice(separatorIndex + separatorLength);
|
||||
|
||||
const dataLines = parseSseEvent(rawEvent);
|
||||
if (dataLines.length === 0) {
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
continue;
|
||||
}
|
||||
|
||||
const dataText = dataLines.join("\n");
|
||||
if (dataText === "[DONE]") {
|
||||
return;
|
||||
}
|
||||
|
||||
const parsed = JSON.parse(dataText) as
|
||||
| OpenAIChatChunk
|
||||
| { error?: { message?: string } };
|
||||
if ("error" in parsed && parsed.error) {
|
||||
throw new Error(parsed.error.message || "Stream error");
|
||||
}
|
||||
yield parsed as OpenAIChatChunk;
|
||||
separatorIndex = buffer.search(/\r?\n\r?\n/);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
import {
|
||||
type LoraModelOption,
|
||||
type ModelOption,
|
||||
ModelSelector,
|
||||
} from "@/components/assistant-ui/model-selector";
|
||||
|
|
@ -28,16 +29,15 @@ import {
|
|||
memo,
|
||||
useCallback,
|
||||
useEffect,
|
||||
useMemo,
|
||||
useRef,
|
||||
useState,
|
||||
} from "react";
|
||||
import {
|
||||
ChatSettingsPanel,
|
||||
type InferenceParams,
|
||||
defaultInferenceParams,
|
||||
} from "./chat-settings-sheet";
|
||||
import { ChatSettingsPanel } from "./chat-settings-sheet";
|
||||
import { db } from "./db";
|
||||
import { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
|
||||
import { ChatRuntimeProvider } from "./runtime-provider";
|
||||
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
import {
|
||||
type CompareHandle,
|
||||
CompareHandlesProvider,
|
||||
|
|
@ -47,47 +47,16 @@ import {
|
|||
import { ThreadSidebar } from "./thread-sidebar";
|
||||
import type { ChatView } from "./types";
|
||||
|
||||
const LORA_MODELS: ModelOption[] = [
|
||||
{
|
||||
id: "outputs/llama-3.1-8b-instruct-lora",
|
||||
name: "meta-llama/Llama-3.1-8B-Instruct",
|
||||
description: "LoRA v1",
|
||||
},
|
||||
{
|
||||
id: "outputs/qwen2.5-7b-lora",
|
||||
name: "Qwen/Qwen2.5-7B-Instruct",
|
||||
description: "LoRA v2",
|
||||
},
|
||||
{
|
||||
id: "outputs/mistral-7b-v0.3-lora",
|
||||
name: "mistralai/Mistral-7B-Instruct-v0.3",
|
||||
description: "LoRA v1",
|
||||
},
|
||||
];
|
||||
|
||||
const GGUF_MODELS: ModelOption[] = [
|
||||
{
|
||||
id: "models/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf",
|
||||
name: "Meta-Llama-3.1-8B-Instruct",
|
||||
description: "Q4_K_M",
|
||||
},
|
||||
{
|
||||
id: "models/Qwen2.5-7B-Instruct-Q5_K_M.gguf",
|
||||
name: "Qwen2.5-7B-Instruct",
|
||||
description: "Q5_K_M",
|
||||
},
|
||||
{
|
||||
id: "models/Mistral-7B-Instruct-v0.3-Q4_K_M.gguf",
|
||||
name: "Mistral-7B-Instruct-v0.3",
|
||||
description: "Q4_K_M",
|
||||
},
|
||||
];
|
||||
|
||||
const SingleContent = memo(function SingleContent({
|
||||
threadId,
|
||||
}: { threadId?: string }): ReactElement {
|
||||
newThreadNonce,
|
||||
}: { threadId?: string; newThreadNonce?: string }): ReactElement {
|
||||
return (
|
||||
<ChatRuntimeProvider modelType="base" initialThreadId={threadId}>
|
||||
<ChatRuntimeProvider
|
||||
modelType="base"
|
||||
initialThreadId={threadId}
|
||||
newThreadNonce={newThreadNonce}
|
||||
>
|
||||
<div className="min-h-0 flex-1">
|
||||
<Thread />
|
||||
</div>
|
||||
|
|
@ -233,28 +202,60 @@ function TopBarActions({
|
|||
}
|
||||
|
||||
export function ChatPage(): ReactElement {
|
||||
const [view, setView] = useState<ChatView>({ mode: "single" });
|
||||
const [view, setView] = useState<ChatView>({
|
||||
mode: "single",
|
||||
newThreadNonce: crypto.randomUUID(),
|
||||
});
|
||||
const [settingsOpen, setSettingsOpen] = useState(false);
|
||||
const [inferenceParams, setInferenceParams] = useState<InferenceParams>(
|
||||
defaultInferenceParams,
|
||||
);
|
||||
const inferenceParams = useChatRuntimeStore((state) => state.params);
|
||||
const setInferenceParams = useChatRuntimeStore((state) => state.setParams);
|
||||
const modelsFromStore = useChatRuntimeStore((state) => state.models);
|
||||
const lorasFromStore = useChatRuntimeStore((state) => state.loras);
|
||||
const modelsError = useChatRuntimeStore((state) => state.modelsError);
|
||||
const { refresh, selectModel, ejectModel } = useChatModelRuntime();
|
||||
|
||||
const handleCheckpointChange = useCallback(
|
||||
(v: string) => setInferenceParams((p) => ({ ...p, checkpoint: v })),
|
||||
(value: string, meta?: { isLora: boolean }) => {
|
||||
void selectModel({ id: value, isLora: meta?.isLora });
|
||||
},
|
||||
[selectModel],
|
||||
);
|
||||
const handleEject = useCallback(() => {
|
||||
void ejectModel();
|
||||
}, [ejectModel]);
|
||||
const handleNewThread = useCallback(
|
||||
() => setView({ mode: "single", newThreadNonce: crypto.randomUUID() }),
|
||||
[],
|
||||
);
|
||||
const handleEject = useCallback(
|
||||
() => setInferenceParams((p) => ({ ...p, checkpoint: "" })),
|
||||
[],
|
||||
);
|
||||
const handleNewThread = useCallback(() => setView({ mode: "single" }), []);
|
||||
const handleNewCompare = useCallback(
|
||||
() => setView({ mode: "compare", pairId: crypto.randomUUID() }),
|
||||
[],
|
||||
);
|
||||
|
||||
const models =
|
||||
inferenceParams.inferenceEngine === "llama-cpp" ? GGUF_MODELS : LORA_MODELS;
|
||||
const models = useMemo<ModelOption[]>(
|
||||
() =>
|
||||
modelsFromStore.map((model) => ({
|
||||
id: model.id,
|
||||
name: model.name,
|
||||
description: model.description,
|
||||
})),
|
||||
[modelsFromStore],
|
||||
);
|
||||
|
||||
const loraModels = useMemo<LoraModelOption[]>(
|
||||
() =>
|
||||
lorasFromStore.map((lora) => ({
|
||||
id: lora.id,
|
||||
name: lora.name,
|
||||
baseModel: lora.baseModel,
|
||||
updatedAt: lora.updatedAt,
|
||||
})),
|
||||
[lorasFromStore],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
void refresh();
|
||||
}, [refresh]);
|
||||
|
||||
return (
|
||||
<SidebarProvider
|
||||
|
|
@ -286,12 +287,18 @@ export function ChatPage(): ReactElement {
|
|||
/>
|
||||
<ModelSelector
|
||||
models={models}
|
||||
loraModels={loraModels}
|
||||
value={inferenceParams.checkpoint}
|
||||
onValueChange={handleCheckpointChange}
|
||||
onEject={handleEject}
|
||||
variant="ghost"
|
||||
/>
|
||||
</div>
|
||||
{modelsError && (
|
||||
<div className="ml-2 text-xs text-destructive truncate max-w-[28rem]">
|
||||
{modelsError}
|
||||
</div>
|
||||
)}
|
||||
<div className="flex-1" />
|
||||
<button
|
||||
type="button"
|
||||
|
|
@ -305,8 +312,9 @@ export function ChatPage(): ReactElement {
|
|||
|
||||
{view.mode === "single" ? (
|
||||
<SingleContent
|
||||
key={view.threadId ?? "new"}
|
||||
key={view.threadId ?? view.newThreadNonce ?? "new"}
|
||||
threadId={view.threadId}
|
||||
newThreadNonce={view.newThreadNonce}
|
||||
/>
|
||||
) : (
|
||||
<CompareContent key={view.pairId} pairId={view.pairId} />
|
||||
|
|
|
|||
|
|
@ -10,7 +10,6 @@ import { Textarea } from "@/components/ui/textarea";
|
|||
import {
|
||||
ArrowDown01Icon,
|
||||
Delete02Icon,
|
||||
EngineIcon,
|
||||
FloppyDiskIcon,
|
||||
PencilEdit01Icon,
|
||||
Settings02Icon,
|
||||
|
|
@ -20,28 +19,13 @@ import { HugeiconsIcon } from "@hugeicons/react";
|
|||
import { AnimatePresence, motion } from "motion/react";
|
||||
import type { ReactNode } from "react";
|
||||
import { useState } from "react";
|
||||
import {
|
||||
DEFAULT_INFERENCE_PARAMS,
|
||||
type InferenceParams,
|
||||
} from "./types/runtime";
|
||||
|
||||
export interface InferenceParams {
|
||||
temperature: number;
|
||||
topP: number;
|
||||
topK: number;
|
||||
repetitionPenalty: number;
|
||||
maxTokens: number;
|
||||
systemPrompt: string;
|
||||
inferenceEngine: string;
|
||||
checkpoint: string;
|
||||
}
|
||||
|
||||
export const defaultInferenceParams: InferenceParams = {
|
||||
temperature: 0.7,
|
||||
topP: 0.9,
|
||||
topK: 50,
|
||||
repetitionPenalty: 1.1,
|
||||
maxTokens: 512,
|
||||
systemPrompt: "",
|
||||
inferenceEngine: "unsloth",
|
||||
checkpoint: "outputs/llama-3.1-8b-instruct-lora",
|
||||
};
|
||||
export const defaultInferenceParams = DEFAULT_INFERENCE_PARAMS;
|
||||
export type { InferenceParams } from "./types/runtime";
|
||||
|
||||
export interface Preset {
|
||||
name: string;
|
||||
|
|
@ -72,11 +56,6 @@ const BUILTIN_PRESETS: Preset[] = [
|
|||
},
|
||||
];
|
||||
|
||||
const ENGINE_OPTIONS = [
|
||||
{ value: "unsloth", label: "Unsloth" },
|
||||
{ value: "llama-cpp", label: "llama.cpp (GGUF)" },
|
||||
];
|
||||
|
||||
function ParamSlider({
|
||||
label,
|
||||
value,
|
||||
|
|
@ -285,33 +264,6 @@ export function ChatSettingsPanel({
|
|||
/>
|
||||
</div>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={EngineIcon}
|
||||
label="Inference Engine"
|
||||
defaultOpen={true}
|
||||
>
|
||||
<div>
|
||||
<span className="mb-1 block text-[11px] text-muted-foreground">
|
||||
Backend
|
||||
</span>
|
||||
<Select
|
||||
value={params.inferenceEngine}
|
||||
onValueChange={set("inferenceEngine")}
|
||||
>
|
||||
<SelectTrigger className="h-8 w-full text-xs corner-squircle">
|
||||
<SelectValue />
|
||||
</SelectTrigger>
|
||||
<SelectContent>
|
||||
{ENGINE_OPTIONS.map((o) => (
|
||||
<SelectItem key={o.value} value={o.value}>
|
||||
{o.label}
|
||||
</SelectItem>
|
||||
))}
|
||||
</SelectContent>
|
||||
</Select>
|
||||
</div>
|
||||
</CollapsibleSection>
|
||||
|
||||
<CollapsibleSection
|
||||
icon={SlidersHorizontalIcon}
|
||||
label="Sampling"
|
||||
|
|
|
|||
|
|
@ -0,0 +1,175 @@
|
|||
import { useCallback } from "react";
|
||||
import { toast } from "sonner";
|
||||
import {
|
||||
getInferenceStatus,
|
||||
listLoras,
|
||||
listModels,
|
||||
loadModel,
|
||||
unloadModel,
|
||||
} from "../api/chat-api";
|
||||
import { useChatRuntimeStore } from "../stores/chat-runtime-store";
|
||||
import type { ChatLoraSummary, ChatModelSummary } from "../types/runtime";
|
||||
|
||||
const DEFAULT_MODEL_MAX_SEQ_LENGTH = 2048;
|
||||
|
||||
type SelectedModelInput = {
|
||||
id: string;
|
||||
isLora?: boolean;
|
||||
};
|
||||
|
||||
const LORA_SUFFIX_RE = /_(\d{9,})$/;
|
||||
|
||||
function parseTrailingEpoch(input: string): number | undefined {
|
||||
const match = input.match(LORA_SUFFIX_RE);
|
||||
if (!match) {
|
||||
return undefined;
|
||||
}
|
||||
const parsed = Number.parseInt(match[1], 10);
|
||||
return Number.isFinite(parsed) ? parsed : undefined;
|
||||
}
|
||||
|
||||
function stripTrailingEpoch(input: string): string {
|
||||
const cleaned = input.replace(LORA_SUFFIX_RE, "").replace(/[_-]+$/, "").trim();
|
||||
return cleaned || input;
|
||||
}
|
||||
|
||||
function describeModel(model: {
|
||||
is_lora?: boolean;
|
||||
is_vision?: boolean;
|
||||
}): string | undefined {
|
||||
const tags: string[] = [];
|
||||
if (model.is_lora) tags.push("LoRA");
|
||||
if (model.is_vision) tags.push("Vision");
|
||||
if (!model.is_lora && !model.is_vision) tags.push("Base");
|
||||
return tags.join(" · ");
|
||||
}
|
||||
|
||||
function toChatModelSummary(model: {
|
||||
id: string;
|
||||
name?: string | null;
|
||||
is_lora?: boolean;
|
||||
is_vision?: boolean;
|
||||
}): ChatModelSummary {
|
||||
return {
|
||||
id: model.id,
|
||||
name: model.name || model.id,
|
||||
description: describeModel(model),
|
||||
isLora: Boolean(model.is_lora),
|
||||
isVision: Boolean(model.is_vision),
|
||||
};
|
||||
}
|
||||
|
||||
function toLoraSummary(lora: {
|
||||
display_name: string;
|
||||
adapter_path: string;
|
||||
base_model?: string | null;
|
||||
}): ChatLoraSummary {
|
||||
const idTail = lora.adapter_path.split("/").filter(Boolean).at(-1) ?? "";
|
||||
const updatedAt =
|
||||
parseTrailingEpoch(lora.display_name) ?? parseTrailingEpoch(idTail);
|
||||
|
||||
return {
|
||||
id: lora.adapter_path,
|
||||
name: stripTrailingEpoch(lora.display_name),
|
||||
baseModel: lora.base_model || "Unknown base model",
|
||||
updatedAt,
|
||||
};
|
||||
}
|
||||
|
||||
export function useChatModelRuntime() {
|
||||
const params = useChatRuntimeStore((state) => state.params);
|
||||
const models = useChatRuntimeStore((state) => state.models);
|
||||
const loras = useChatRuntimeStore((state) => state.loras);
|
||||
const setModels = useChatRuntimeStore((state) => state.setModels);
|
||||
const setLoras = useChatRuntimeStore((state) => state.setLoras);
|
||||
const setModelsError = useChatRuntimeStore((state) => state.setModelsError);
|
||||
const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint);
|
||||
const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint);
|
||||
|
||||
const refresh = useCallback(async () => {
|
||||
setModelsError(null);
|
||||
try {
|
||||
const [listRes, statusRes, lorasRes] = await Promise.all([
|
||||
listModels(),
|
||||
getInferenceStatus(),
|
||||
listLoras(),
|
||||
]);
|
||||
|
||||
setModels(listRes.models.map(toChatModelSummary));
|
||||
setLoras(lorasRes.loras.map(toLoraSummary));
|
||||
|
||||
if (statusRes.active_model) {
|
||||
setCheckpoint(statusRes.active_model);
|
||||
}
|
||||
} catch (error) {
|
||||
const message =
|
||||
error instanceof Error ? error.message : "Failed to load models";
|
||||
setModelsError(message);
|
||||
}
|
||||
}, [setCheckpoint, setLoras, setModels, setModelsError]);
|
||||
|
||||
const selectModel = useCallback(
|
||||
async (selection: string | SelectedModelInput) => {
|
||||
const modelId = typeof selection === "string" ? selection : selection.id;
|
||||
if (!modelId || params.checkpoint === modelId) {
|
||||
return;
|
||||
}
|
||||
|
||||
const explicitIsLora =
|
||||
typeof selection === "string" ? undefined : selection.isLora;
|
||||
const model = models.find((entry) => entry.id === modelId);
|
||||
const lora = loras.find((entry) => entry.id === modelId);
|
||||
const isLora =
|
||||
explicitIsLora ?? model?.isLora ?? (lora ? true : false);
|
||||
const displayName = model?.name || lora?.name || modelId;
|
||||
const loadingToastId = toast.loading(`Loading ${displayName}...`);
|
||||
|
||||
setModelsError(null);
|
||||
try {
|
||||
if (params.checkpoint) {
|
||||
await unloadModel({ model_path: params.checkpoint });
|
||||
}
|
||||
|
||||
await loadModel({
|
||||
model_path: modelId,
|
||||
hf_token: null,
|
||||
max_seq_length: DEFAULT_MODEL_MAX_SEQ_LENGTH,
|
||||
load_in_4bit: true,
|
||||
is_lora: isLora,
|
||||
});
|
||||
|
||||
setCheckpoint(modelId);
|
||||
await refresh();
|
||||
toast.success(`${displayName} loaded`, { id: loadingToastId });
|
||||
} catch (error) {
|
||||
const message =
|
||||
error instanceof Error ? error.message : "Failed to load model";
|
||||
setModelsError(message);
|
||||
toast.error(message, { id: loadingToastId });
|
||||
}
|
||||
},
|
||||
[loras, models, params.checkpoint, refresh, setCheckpoint, setModelsError],
|
||||
);
|
||||
|
||||
const ejectModel = useCallback(async () => {
|
||||
if (!params.checkpoint) {
|
||||
return;
|
||||
}
|
||||
setModelsError(null);
|
||||
try {
|
||||
await unloadModel({ model_path: params.checkpoint });
|
||||
clearCheckpoint();
|
||||
await refresh();
|
||||
} catch (error) {
|
||||
const message =
|
||||
error instanceof Error ? error.message : "Failed to unload model";
|
||||
setModelsError(message);
|
||||
}
|
||||
}, [clearCheckpoint, params.checkpoint, refresh, setModelsError]);
|
||||
|
||||
return {
|
||||
refresh,
|
||||
selectModel,
|
||||
ejectModel,
|
||||
};
|
||||
}
|
||||
|
|
@ -5,3 +5,5 @@ export {
|
|||
type InferenceParams,
|
||||
type Preset,
|
||||
} from "./chat-settings-sheet";
|
||||
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
|
||||
export { useChatModelRuntime } from "./hooks/use-chat-model-runtime";
|
||||
|
|
|
|||
|
|
@ -7,9 +7,9 @@ import {
|
|||
type ExportedMessageRepositoryItem,
|
||||
type PendingAttachment,
|
||||
RuntimeAdapterProvider,
|
||||
Suggestions,
|
||||
SimpleImageAttachmentAdapter,
|
||||
SimpleTextAttachmentAdapter,
|
||||
Suggestions,
|
||||
type ThreadHistoryAdapter,
|
||||
type ThreadMessage,
|
||||
type ThreadUserMessagePart,
|
||||
|
|
@ -24,10 +24,17 @@ import { createAssistantStream } from "assistant-stream";
|
|||
import mammoth from "mammoth";
|
||||
import { type ReactElement, type ReactNode, useEffect, useMemo } from "react";
|
||||
import { extractText, getDocumentProxy } from "unpdf";
|
||||
import { createStreamAdapter } from "./adapter";
|
||||
import { createOpenAIStreamAdapter } from "./api/chat-adapter";
|
||||
import { db } from "./db";
|
||||
import type { MessageRecord, ModelType } from "./types";
|
||||
|
||||
const DEFAULT_SUGGESTIONS = [
|
||||
"Draw a simple flowchart of a login system using Mermaid",
|
||||
"Solve the integral of x²·sin(x) step by step",
|
||||
"Write a Python function that finds the longest palindrome in a string",
|
||||
"Format a comparison of 3 databases as a markdown table with pros and cons",
|
||||
];
|
||||
|
||||
class PDFAttachmentAdapter implements AttachmentAdapter {
|
||||
accept = "application/pdf";
|
||||
|
||||
|
|
@ -288,9 +295,11 @@ function ThreadHistoryProvider({
|
|||
);
|
||||
}
|
||||
|
||||
const chatAdapter = createStreamAdapter();
|
||||
const useRuntimeHook = (): ReturnType<typeof useLocalRuntime> =>
|
||||
useLocalRuntime(chatAdapter);
|
||||
const chatAdapter = createOpenAIStreamAdapter();
|
||||
|
||||
function useRuntimeHook(): ReturnType<typeof useLocalRuntime> {
|
||||
return useLocalRuntime(chatAdapter);
|
||||
}
|
||||
|
||||
function ThreadAutoSwitch({
|
||||
threadId,
|
||||
|
|
@ -308,16 +317,33 @@ function ThreadAutoSwitch({
|
|||
return null;
|
||||
}
|
||||
|
||||
function ThreadNewChatSwitch({
|
||||
nonce,
|
||||
}: { nonce: string }): ReactElement | null {
|
||||
const aui = useAui();
|
||||
const isLoading = useAuiState(({ threads }) => threads.isLoading);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isLoading) {
|
||||
aui.threads().switchToNewThread();
|
||||
}
|
||||
}, [aui, isLoading, nonce]);
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
export function ChatRuntimeProvider({
|
||||
children,
|
||||
modelType = "base",
|
||||
pairId,
|
||||
initialThreadId,
|
||||
newThreadNonce,
|
||||
}: {
|
||||
children: ReactNode;
|
||||
modelType?: ModelType;
|
||||
pairId?: string;
|
||||
initialThreadId?: string;
|
||||
newThreadNonce?: string;
|
||||
}): ReactElement {
|
||||
const runtime = useRemoteThreadListRuntime({
|
||||
runtimeHook: useRuntimeHook,
|
||||
|
|
@ -328,17 +354,15 @@ export function ChatRuntimeProvider({
|
|||
});
|
||||
|
||||
const aui = useAui({
|
||||
suggestions: Suggestions([
|
||||
"Draw a simple flowchart of a login system using Mermaid",
|
||||
"Solve the integral of x\u00B2\u00B7sin(x) step by step",
|
||||
"Write a Python function that finds the longest palindrome in a string",
|
||||
"Format a comparison of 3 databases as a markdown table with pros and cons",
|
||||
]),
|
||||
suggestions: Suggestions(DEFAULT_SUGGESTIONS),
|
||||
});
|
||||
|
||||
return (
|
||||
<AssistantRuntimeProvider runtime={runtime} aui={aui}>
|
||||
{initialThreadId && <ThreadAutoSwitch threadId={initialThreadId} />}
|
||||
{!initialThreadId && newThreadNonce && (
|
||||
<ThreadNewChatSwitch nonce={newThreadNonce} />
|
||||
)}
|
||||
{children}
|
||||
</AssistantRuntimeProvider>
|
||||
);
|
||||
|
|
|
|||
|
|
@ -0,0 +1,58 @@
|
|||
import { create } from "zustand";
|
||||
import {
|
||||
DEFAULT_INFERENCE_PARAMS,
|
||||
type ChatLoraSummary,
|
||||
type ChatModelSummary,
|
||||
type InferenceParams,
|
||||
} from "../types/runtime";
|
||||
|
||||
type ChatRuntimeStore = {
|
||||
params: InferenceParams;
|
||||
models: ChatModelSummary[];
|
||||
loras: ChatLoraSummary[];
|
||||
warmingByThreadId: Record<string, boolean>;
|
||||
modelsError: string | null;
|
||||
setParams: (params: InferenceParams) => void;
|
||||
setModels: (models: ChatModelSummary[]) => void;
|
||||
setLoras: (loras: ChatLoraSummary[]) => void;
|
||||
setThreadWarming: (threadId: string, warming: boolean) => void;
|
||||
setModelsError: (error: string | null) => void;
|
||||
setCheckpoint: (modelId: string) => void;
|
||||
clearCheckpoint: () => void;
|
||||
};
|
||||
|
||||
export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
|
||||
params: DEFAULT_INFERENCE_PARAMS,
|
||||
models: [],
|
||||
loras: [],
|
||||
warmingByThreadId: {},
|
||||
modelsError: null,
|
||||
setParams: (params) => set({ params }),
|
||||
setModels: (models) => set({ models }),
|
||||
setLoras: (loras) => set({ loras }),
|
||||
setThreadWarming: (threadId, warming) =>
|
||||
set((state) => {
|
||||
const next = { ...state.warmingByThreadId };
|
||||
if (warming) {
|
||||
next[threadId] = true;
|
||||
} else {
|
||||
delete next[threadId];
|
||||
}
|
||||
return { warmingByThreadId: next };
|
||||
}),
|
||||
setModelsError: (modelsError) => set({ modelsError }),
|
||||
setCheckpoint: (modelId) =>
|
||||
set((state) => ({
|
||||
params: {
|
||||
...state.params,
|
||||
checkpoint: modelId,
|
||||
},
|
||||
})),
|
||||
clearCheckpoint: () =>
|
||||
set((state) => ({
|
||||
params: {
|
||||
...state.params,
|
||||
checkpoint: "",
|
||||
},
|
||||
})),
|
||||
}));
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
export type ModelType = "base" | "lora";
|
||||
|
||||
export type ChatView =
|
||||
| { mode: "single"; threadId?: string }
|
||||
| { mode: "single"; threadId?: string; newThreadNonce?: string }
|
||||
| { mode: "compare"; pairId: string };
|
||||
|
||||
export interface ThreadRecord {
|
||||
|
|
|
|||
79
studio/frontend/src/features/chat/types/api.ts
Normal file
79
studio/frontend/src/features/chat/types/api.ts
Normal file
|
|
@ -0,0 +1,79 @@
|
|||
export interface BackendModelDetails {
|
||||
id: string;
|
||||
name?: string | null;
|
||||
is_vision?: boolean;
|
||||
is_lora?: boolean;
|
||||
}
|
||||
|
||||
export interface ListModelsResponse {
|
||||
models: BackendModelDetails[];
|
||||
default_models: string[];
|
||||
}
|
||||
|
||||
export interface BackendLoraInfo {
|
||||
display_name: string;
|
||||
adapter_path: string;
|
||||
base_model?: string | null;
|
||||
}
|
||||
|
||||
export interface ListLorasResponse {
|
||||
loras: BackendLoraInfo[];
|
||||
outputs_dir: string;
|
||||
}
|
||||
|
||||
export interface LoadModelRequest {
|
||||
model_path: string;
|
||||
hf_token: string | null;
|
||||
max_seq_length: number;
|
||||
load_in_4bit: boolean;
|
||||
is_lora: boolean;
|
||||
}
|
||||
|
||||
export interface LoadModelResponse {
|
||||
status: string;
|
||||
model: string;
|
||||
display_name: string;
|
||||
is_vision: boolean;
|
||||
is_lora: boolean;
|
||||
}
|
||||
|
||||
export interface UnloadModelRequest {
|
||||
model_path: string;
|
||||
}
|
||||
|
||||
export interface InferenceStatusResponse {
|
||||
active_model: string | null;
|
||||
is_vision: boolean;
|
||||
loading: string[];
|
||||
loaded: string[];
|
||||
}
|
||||
|
||||
export interface OpenAIChatMessage {
|
||||
role: "system" | "user" | "assistant";
|
||||
content: string;
|
||||
}
|
||||
|
||||
export interface OpenAIChatCompletionsRequest {
|
||||
model: string;
|
||||
messages: OpenAIChatMessage[];
|
||||
stream: boolean;
|
||||
temperature: number;
|
||||
top_p: number;
|
||||
max_tokens: number;
|
||||
top_k: number;
|
||||
repetition_penalty: number;
|
||||
}
|
||||
|
||||
export interface OpenAIChatDelta {
|
||||
role?: string;
|
||||
content?: string;
|
||||
}
|
||||
|
||||
export interface OpenAIChatChunkChoice {
|
||||
delta?: OpenAIChatDelta;
|
||||
finish_reason?: string | null;
|
||||
}
|
||||
|
||||
export interface OpenAIChatChunk {
|
||||
choices?: OpenAIChatChunkChoice[];
|
||||
}
|
||||
34
studio/frontend/src/features/chat/types/runtime.ts
Normal file
34
studio/frontend/src/features/chat/types/runtime.ts
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
export interface InferenceParams {
|
||||
temperature: number;
|
||||
topP: number;
|
||||
topK: number;
|
||||
repetitionPenalty: number;
|
||||
maxTokens: number;
|
||||
systemPrompt: string;
|
||||
checkpoint: string;
|
||||
}
|
||||
|
||||
export const DEFAULT_INFERENCE_PARAMS: InferenceParams = {
|
||||
temperature: 0.7,
|
||||
topP: 0.9,
|
||||
topK: 50,
|
||||
repetitionPenalty: 1.1,
|
||||
maxTokens: 512,
|
||||
systemPrompt: "",
|
||||
checkpoint: "",
|
||||
};
|
||||
|
||||
export interface ChatModelSummary {
|
||||
id: string;
|
||||
name: string;
|
||||
description?: string;
|
||||
isVision: boolean;
|
||||
isLora: boolean;
|
||||
}
|
||||
|
||||
export interface ChatLoraSummary {
|
||||
id: string;
|
||||
name: string;
|
||||
baseModel: string;
|
||||
updatedAt?: number;
|
||||
}
|
||||
|
|
@ -0,0 +1,54 @@
|
|||
import type { ChatModelRunResult } from "@assistant-ui/react";
|
||||
|
||||
type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
|
||||
|
||||
const THINK_OPEN_TAG = "<think>";
|
||||
const THINK_CLOSE_TAG = "</think>";
|
||||
|
||||
function appendTextPart(parts: ContentPart[], text: string): void {
|
||||
if (text) {
|
||||
parts.push({ type: "text", text });
|
||||
}
|
||||
}
|
||||
|
||||
function appendReasoningPart(parts: ContentPart[], text: string): void {
|
||||
if (text) {
|
||||
parts.push({ type: "reasoning", text });
|
||||
}
|
||||
}
|
||||
|
||||
export function parseAssistantContent(
|
||||
raw: string,
|
||||
): ContentPart[] {
|
||||
const parts: ContentPart[] = [];
|
||||
if (!raw) {
|
||||
return parts;
|
||||
}
|
||||
|
||||
let cursor = 0;
|
||||
while (cursor < raw.length) {
|
||||
const openIndex = raw.indexOf(THINK_OPEN_TAG, cursor);
|
||||
if (openIndex === -1) {
|
||||
appendTextPart(parts, raw.slice(cursor));
|
||||
break;
|
||||
}
|
||||
|
||||
appendTextPart(parts, raw.slice(cursor, openIndex));
|
||||
|
||||
const reasoningStart = openIndex + THINK_OPEN_TAG.length;
|
||||
const closeIndex = raw.indexOf(THINK_CLOSE_TAG, reasoningStart);
|
||||
if (closeIndex === -1) {
|
||||
appendReasoningPart(parts, raw.slice(reasoningStart));
|
||||
break;
|
||||
}
|
||||
|
||||
appendReasoningPart(parts, raw.slice(reasoningStart, closeIndex));
|
||||
cursor = closeIndex + THINK_CLOSE_TAG.length;
|
||||
}
|
||||
|
||||
return parts;
|
||||
}
|
||||
|
||||
export function hasClosedThinkTag(raw: string): boolean {
|
||||
return raw.includes(THINK_CLOSE_TAG);
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue