feat: support muxing fim and chat (stacklok#343)

* chore: update deps

* chore: manually update openapi and generate new types

* feat: add muxing fim, chat support

* refactor: rename provider delete text

* refactor: muxing based on new enum

* chore: update enum

Author: peppescg

package-lock.json

@@ -31,7 +31,7 @@
     "@radix-ui/react-dialog": "^1.1.4",
     "@radix-ui/react-separator": "^1.1.0",
     "@radix-ui/react-slot": "^1.1.0",
-    "@stacklok/ui-kit": "^1.0.1-4",
+    "@stacklok/ui-kit": "^1.0.1-9",
     "@tanstack/react-query": "^5.64.1",
     "@tanstack/react-query-devtools": "^5.66.0",
     "@types/lodash": "^4.17.15",

package.json

@@ -144,11 +144,23 @@ export type ModelByProvider = {
 
 /**
  * Represents the different types of matchers we support.
+ *
+ * The 3 rules present match filenames and request types. They're used in conjunction with the
+ * matcher field in the MuxRule model.
+ * E.g.
+ * - catch_all-> Always match
+ * - filename_match and match: requests.py -> Match the request if the filename is requests.py
+ * - fim_filename and match: main.py -> Match the request if the request type is fim
+ * and the filename is main.py
+ *
+ * NOTE: Removing or updating fields from this enum will require a migration.
+ * Adding new fields is safe.
  */
 export enum MuxMatcherType {
   CATCH_ALL = 'catch_all',
   FILENAME_MATCH = 'filename_match',
-  REQUEST_TYPE_MATCH = 'request_type_match',
+  FIM_FILENAME = 'fim_filename',
+  CHAT_FILENAME = 'chat_filename',
 }
 
 /**

src/api/generated/types.gen.ts

@@ -1642,10 +1642,11 @@
         "enum": [
           "catch_all",
           "filename_match",
-          "request_type_match"
+          "fim_filename",
+          "chat_filename"
         ],
         "title": "MuxMatcherType",
-        "description": "Represents the different types of matchers we support."
+        "description": "Represents the different types of matchers we support.\n\nThe 3 rules present match filenames and request types. They're used in conjunction with the\nmatcher field in the MuxRule model.\nE.g.\n- catch_all-> Always match\n- filename_match and match: requests.py -> Match the request if the filename is requests.py\n- fim_filename and match: main.py -> Match the request if the request type is fim\nand the filename is main.py\n\nNOTE: Removing or updating fields from this enum will require a migration.\nAdding new fields is safe."
       },
       "MuxRule": {
         "properties": {

src/api/openapi.json

@@ -10,7 +10,10 @@ export function WorkspacesByProvider({
   if (workspaces.length === 0) return null
   return (
     <div className="mb-6 flex flex-col gap-1">
-      <p>The following workspaces will be impacted by this action</p>
+      <p>
+        The following workspaces are currently using this provider and will need
+        to be updated:
+      </p>
       <div className="flex flex-wrap gap-1">
         {uniqBy(workspaces, 'name').map((item, index) => {
           return (

src/features/providers/components/workspaces-by-provider.tsx

@@ -7,6 +7,12 @@ test('renders muxing model', async () => {
   render(
     <WorkspaceMuxingModel isArchived={false} workspaceName="fake-workspace" />
   )
+
+  expect(
+    screen.getByRole('button', {
+      name: /all types/i,
+    })
+  ).toBeVisible()
   expect(screen.getByText(/model muxing/i)).toBeVisible()
   expect(
     screen.getByText(
@@ -46,6 +52,11 @@ test('disabled muxing fields and buttons for archived workspace', async () => {
 
   expect(await screen.findByRole('button', { name: /save/i })).toBeDisabled()
   expect(screen.getByTestId(/workspace-models-dropdown/i)).toBeDisabled()
+  expect(
+    screen.getByRole('button', {
+      name: /all types/i,
+    })
+  ).toBeDisabled()
   expect(
     await screen.findByRole('button', { name: /add filter/i })
   ).toBeDisabled()
@@ -75,12 +86,27 @@ test('submit additional model overrides', async () => {
     name: /filter by/i,
   })
   expect(textFields.length).toEqual(2)
+
+  const requestTypeSelect = screen.getAllByRole('button', {
+    name: /fim & chat/i,
+  })[0]
+  await userEvent.click(requestTypeSelect as HTMLFormElement)
+  await userEvent.click(
+    screen.getByRole('option', {
+      name: 'FIM',
+    })
+  )
+  expect(
+    screen.getByRole('button', {
+      name: 'FIM',
+    })
+  ).toBeVisible()
   const modelsButton = await screen.findAllByTestId(
     /workspace-models-dropdown/i
   )
   expect(modelsButton.length).toEqual(2)
 
-  await userEvent.type(textFields[1] as HTMLFormElement, '.ts')
+  await userEvent.type(textFields[0] as HTMLFormElement, '.tsx')
 
   await userEvent.click(
     (await screen.findByRole('button', {
@@ -94,6 +120,37 @@ test('submit additional model overrides', async () => {
     })
   )
 
+  await userEvent.click(screen.getByRole('button', { name: /add filter/i }))
+  await userEvent.click(
+    screen.getAllByRole('button', {
+      name: /chat/i,
+    })[1] as HTMLFormElement
+  )
+
+  await userEvent.click(
+    screen.getByRole('option', {
+      name: 'Chat',
+    })
+  )
+
+  await userEvent.type(
+    screen.getAllByRole('textbox', {
+      name: /filter by/i,
+    })[1] as HTMLFormElement,
+    '.ts'
+  )
+
+  await userEvent.click(
+    (await screen.findByRole('button', {
+      name: /select a model/i,
+    })) as HTMLFormElement
+  )
+
+  await userEvent.click(
+    screen.getByRole('option', {
+      name: /chatgpt-4o/i,
+    })
+  )
   await userEvent.click(screen.getByRole('button', { name: /save/i }))
 
   await waitFor(() => {

src/features/workspace/components/__tests__/workspace-muxing-model.test.tsx

@@ -37,7 +37,7 @@ function groupModelsByProviderName(
     id: providerName,
     textValue: providerName,
     items: items.map((item) => ({
-      id: `${item.provider_id}/${item.name}`,
+      id: `${item.provider_id}:${item.name}`,
       textValue: item.name,
     })),
   }))
@@ -116,7 +116,7 @@ export function WorkspaceModelsDropdown({
               const selectedValue = v.values().next().value
               if (!selectedValue && typeof selectedValue !== 'string') return
               if (typeof selectedValue === 'string') {
-                const [provider_id, modelName] = selectedValue.split('/')
+                const [provider_id, modelName] = selectedValue.split(':')
                 if (!provider_id || !modelName) return
                 onChange({
                   model: modelName,

src/features/workspace/components/workspace-models-dropdown.tsx

@@ -9,6 +9,8 @@ import {
   Label,
   Link,
   LinkButton,
+  Select,
+  SelectButton,
   Text,
   TextField,
   Tooltip,
@@ -17,10 +19,7 @@ import {
 } from '@stacklok/ui-kit'
 import { twMerge } from 'tailwind-merge'
 import { useMutationPreferredModelWorkspace } from '../hooks/use-mutation-preferred-model-workspace'
-import {
-  MuxMatcherType,
-  V1ListAllModelsForAllProvidersResponse,
-} from '@/api/generated'
+import { V1ListAllModelsForAllProvidersResponse } from '@/api/generated'
 import { FormEvent } from 'react'
 import {
   LayersThree01,
@@ -37,6 +36,7 @@ import {
   useMuxingRulesFormState,
 } from '../hooks/use-muxing-rules-form-workspace'
 import { FormButtons } from '@/components/FormButtons'
+import { getRuleData, isRequestType } from '../lib/utils'
 
 function MissingProviderBanner() {
   return (
@@ -77,9 +77,31 @@ function SortableItem({
   isArchived,
   isDefaultRule,
 }: SortableItemProps) {
-  const placeholder = isDefaultRule ? 'Catch-all' : 'e.g. file type, file name'
+  const { selectedKey, placeholder, items } = getRuleData({
+    isDefaultRule,
+    matcher_type: rule.matcher_type,
+  })
+
   return (
     <div className="flex items-center gap-2" key={rule.id}>
+      <div className="flex w-2/5 justify-between">
+        <Select
+          aria-labelledby="request type"
+          selectedKey={selectedKey}
+          name="request_type"
+          isRequired
+          isDisabled={isDefaultRule}
+          className="w-full"
+          items={items}
+          onSelectionChange={(matcher_type) => {
+            if (isRequestType(matcher_type)) {
+              setRuleItem({ ...rule, matcher_type })
+            }
+          }}
+        >
+          <SelectButton />
+        </Select>
+      </div>
       <div className="flex w-full justify-between">
         <TextField
           aria-labelledby="filter-by-label-id"
@@ -149,9 +171,7 @@ export function WorkspaceMuxingModel({
         body: rules.map(({ id, ...rest }) => {
           void id
 
-          return rest.matcher
-            ? { ...rest, matcher_type: MuxMatcherType.FILENAME_MATCH }
-            : { ...rest }
+          return rest
         }),
       },
       {
@@ -200,6 +220,7 @@ export function WorkspaceMuxingModel({
           <div className="flex w-full flex-col gap-2">
             <div className="flex gap-2">
               <div className="w-12">&nbsp;</div>
+              <div className="w-2/5">Request Type</div>
               <div className="w-full">
                 <Label id="filter-by-label-id" className="flex items-center">
                   Filter by
@@ -214,7 +235,7 @@ export function WorkspaceMuxingModel({
                 </Label>
               </div>
               <div className="w-3/5">
-                <Label id="preferred-model-id">Preferred Model</Label>
+                <Label id="preferred-model-id">Model</Label>
               </div>
             </div>
             <SortableArea