UI feat: add usage field to model form

Signed-off-by: Ryan Hopper-Lowe <[email protected]>

ui/admin/app/components/agent/AgentForm.tsx

@@ -1,6 +1,6 @@
 import { zodResolver } from "@hookform/resolvers/zod";
 import { BrainIcon } from "lucide-react";
-import { useEffect } from "react";
+import { useEffect, useMemo } from "react";
 import { useForm } from "react-hook-form";
 import useSWR from "swr";
 import { z } from "zod";
@@ -41,11 +41,17 @@ type AgentFormProps = {
 };
 
 export function AgentForm({ agent, onSubmit, onChange }: AgentFormProps) {
-    const { data: models } = useSWR(
+    const getModels = useSWR(
         ModelApiService.getModels.key(),
         ModelApiService.getModels
     );
 
+    const models = useMemo(() => {
+        if (!getModels.data) return [];
+
+        return getModels.data.filter((m) => !m.usage || m.usage === "agent");
+    }, [getModels.data]);
+
     const form = useForm<AgentInfoFormValues>({
         resolver: zodResolver(formSchema),
         mode: "onChange",
@@ -125,15 +131,12 @@ export function AgentForm({ agent, onSubmit, onChange }: AgentFormProps) {
                                 <SelectEmptyItem>
                                     Use System Default
                                 </SelectEmptyItem>
-                                {models
-                                    ?.filter(
-                                        (m) => !m.usage || m.usage === "agent"
-                                    )
-                                    .map((m) => (
-                                        <SelectItem key={m.id} value={m.id}>
-                                            {m.name || m.id}
-                                        </SelectItem>
-                                    ))}
+
+                                {models.map((m) => (
+                                    <SelectItem key={m.id} value={m.id}>
+                                        {m.name || m.id}
+                                    </SelectItem>
+                                ))}
                             </SelectContent>
                         </Select>
                     )}

ui/admin/app/components/model/ModelForm.tsx

@@ -10,6 +10,8 @@ import {
     ModelManifest,
     ModelManifestSchema,
     ModelProvider,
+    ModelUsage,
+    getModelUsageLabel,
     getModelsForProvider,
 } from "~/lib/model/models";
 import { BadRequestError } from "~/lib/service/api/apiErrors";
@@ -70,6 +72,7 @@ export function ModelForm(props: ModelFormProps) {
             modelProvider: model?.modelProvider ?? "",
             active: model?.active ?? true,
             default: model?.default ?? false,
+            usage: model?.usage ?? "agent",
         };
     }, [model]);
 
@@ -157,6 +160,30 @@ export function ModelForm(props: ModelFormProps) {
                     }}
                 </ControlledCustomInput>
 
+                <ControlledCustomInput
+                    control={form.control}
+                    name="usage"
+                    label="Usage"
+                >
+                    {({ field: { ref: _, ...field }, className }) => (
+                        <Select {...field} onValueChange={field.onChange}>
+                            <SelectTrigger className={className}>
+                                <SelectValue placeholder="Select Usage" />
+                            </SelectTrigger>
+
+                            <SelectContent>
+                                {Object.entries(ModelUsage).map(
+                                    ([key, value]) => (
+                                        <SelectItem key={key} value={value}>
+                                            {getModelUsageLabel(value)}
+                                        </SelectItem>
+                                    )
+                                )}
+                            </SelectContent>
+                        </Select>
+                    )}
+                </ControlledCustomInput>
+
                 <ControlledCheckbox
                     control={form.control}
                     name="default"

ui/admin/app/lib/model/models.ts

@@ -2,13 +2,32 @@ import { z } from "zod";
 
 import { EntityMeta } from "~/lib/model/primitives";
 
+export const ModelUsage = {
+    Agent: "agent",
+    TextEmbedding: "text-embedding",
+    ImageGeneration: "image-generation",
+} as const;
+export type ModelUsage = (typeof ModelUsage)[keyof typeof ModelUsage];
+
+const ModelUsageLabels = {
+    [ModelUsage.Agent]: "Agent",
+    [ModelUsage.TextEmbedding]: "Text Embedding",
+    [ModelUsage.ImageGeneration]: "Image Generation",
+} as const;
+
+export const getModelUsageLabel = (usage: string) => {
+    if (!(usage in ModelUsageLabels)) return usage;
+
+    return ModelUsageLabels[usage as ModelUsage];
+};
+
 export type ModelManifest = {
     name?: string;
     targetModel?: string;
     modelProvider: string;
     active: boolean;
     default: boolean;
-    usage?: string;
+    usage: ModelUsage;
 };
 
 export type ModelProviderStatus = {
@@ -24,6 +43,7 @@ export const ModelManifestSchema = z.object({
     modelProvider: z.string().min(1, "Required"),
     active: z.boolean(),
     default: z.boolean(),
+    usage: z.nativeEnum(ModelUsage),
 });
 
 export type ModelProvider = EntityMeta & {