Absolutely glorious domain invariant enforcement

Author: Veikkosuhonen

src/client/components/ChatV2/ChatV2.tsx

@@ -7,17 +7,15 @@ import { enqueueSnackbar } from 'notistack'
 import { useCallback, useEffect, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { useParams, useSearchParams } from 'react-router-dom'
-import { DEFAULT_MODEL, DEFAULT_MODEL_TEMPERATURE, FREE_MODEL, inProduction, validModels } from '../../../config'
+import { DEFAULT_MODEL, DEFAULT_MODEL_TEMPERATURE, FREE_MODEL, inProduction, ValidModelName, ValidModelNameSchema, validModels } from '../../../config'
 import type { ChatMessage, MessageGenerationInfo, ToolCallResultEvent } from '../../../shared/chat'
-import type { RagIndexAttributes } from '../../../shared/types'
 import { getLanguageValue } from '../../../shared/utils'
 import { useIsEmbedded } from '../../contexts/EmbeddedContext'
 import { useChatScroll } from './useChatScroll'
 import useCourse from '../../hooks/useCourse'
 import useCurrentUser from '../../hooks/useCurrentUser'
 import useInfoTexts from '../../hooks/useInfoTexts'
 import useLocalStorageState from '../../hooks/useLocalStorageState'
-import { useCourseRagIndices } from '../../hooks/useRagIndices'
 import useRetryTimeout from '../../hooks/useRetryTimeout'
 import useUserStatus from '../../hooks/useUserStatus'
 import { useAnalyticsDispatch } from '../../stores/analytics'
@@ -31,31 +29,37 @@ import { handleCompletionStreamError } from './error'
 import ToolResult from './ToolResult'
 import { OutlineButtonBlack } from './general/Buttons'
 import { ChatInfo } from './general/ChatInfo'
-import RagSelector, { RagSelectorDescription } from './RagSelector'
 import { SettingsModal } from './SettingsModal'
 import { useChatStream } from './useChatStream'
 import { postCompletionStreamV3 } from './api'
 import PromptSelector from './PromptSelector'
 import ModelSelector from './ModelSelector'
 import { ConversationSplash } from './general/ConversationSplash'
 import { PromptStateProvider, usePromptState } from './PromptState'
+import z from 'zod/v4'
 
-function useLocalStorageStateWithURLDefault(key: string, defaultValue: string, urlKey: string) {
+function useLocalStorageStateWithURLDefault<T>(key: string, defaultValue: string, urlKey: string, schema: z.ZodType<T>) {
   const [value, setValue] = useLocalStorageState(key, defaultValue)
   const [searchParams, setSearchParams] = useSearchParams()
   const urlValue = searchParams.get(urlKey)
 
   // If urlValue is defined, it overrides the localStorage setting.
   // However if user changes the setting, the urlValue is removed.
-  const modifiedSetValue = (newValue: string) => {
+  const modifiedSetValue = (newValue: T) => {
     if (newValue !== urlValue) {
-      setValue(newValue)
+      if (typeof newValue === 'string') {
+        setValue(newValue)
+      } else {
+        setValue(String(newValue))
+      }
       searchParams.delete(urlKey)
       setSearchParams(searchParams)
     }
   }
 
-  return [urlValue ?? value, modifiedSetValue] as const
+  const parsedValue = schema.parse(urlValue ?? value)
+
+  return [parsedValue, modifiedSetValue] as const
 }
 
 const ChatV2Content = () => {
@@ -65,7 +69,6 @@ const ChatV2Content = () => {
   const isMobile = useMediaQuery(theme.breakpoints.down('md'))
 
   const { data: course } = useCourse(courseId)
-  const { ragIndices } = useCourseRagIndices(course?.id)
   const { infoTexts } = useInfoTexts()
 
   const { userStatus, isLoading: statusLoading, refetch: refetchStatus } = useUserStatus(courseId)
@@ -74,12 +77,13 @@ const ChatV2Content = () => {
 
   // local storage states
   const localStoragePrefix = courseId ? `course-${courseId}` : 'general'
-  const [activeModel, setActiveModel] = useLocalStorageStateWithURLDefault('model-v2', DEFAULT_MODEL, 'model')
+  const [activeModel, setActiveModel] = useLocalStorageStateWithURLDefault('model-v2', DEFAULT_MODEL, 'model', ValidModelNameSchema)
   const [disclaimerStatus, setDisclaimerStatus] = useLocalStorageState<boolean>('disclaimer-status', true)
   const [modelTemperature, setModelTemperature] = useLocalStorageStateWithURLDefault(
     `${localStoragePrefix}-chat-model-temperature`,
     String(DEFAULT_MODEL_TEMPERATURE),
     'temperature',
+    z.number(),
   )
 
   const [messages, setMessages] = useLocalStorageState(`${localStoragePrefix}-chat-messages`, [] as ChatMessage[])
@@ -90,12 +94,9 @@ const ChatV2Content = () => {
   const [fileName, setFileName] = useState<string>('')
   const [tokenUsageWarning, setTokenUsageWarning] = useState<string>('')
   const [tokenUsageAlertOpen, setTokenUsageAlertOpen] = useState<boolean>(false)
-  const [allowedModels, setAllowedModels] = useState<string[]>([])
+  const [allowedModels, setAllowedModels] = useState<ValidModelName[]>([])
   const [chatLeftSidePanelOpen, setChatLeftSidePanelOpen] = useState<boolean>(false)
-  // RAG states
-  const [ragIndexId, _setRagIndexId] = useState<number | undefined>()
   const [activeToolResult, setActiveToolResult0] = useState<ToolCallResultEvent | undefined>()
-  const ragIndex = ragIndices?.find((index) => index.id === ragIndexId)
 
   // Analytics
   const dispatchAnalytics = useAnalyticsDispatch()
@@ -106,11 +107,9 @@ const ChatV2Content = () => {
         model: activeModel,
         courseId,
         nMessages: messages.length,
-        ragIndexId,
-        ragIndexName: ragIndex?.metadata.name,
       },
     })
-  }, [messages, courseId, ragIndexId, activeModel, dispatchAnalytics])
+  }, [messages, courseId, activeModel, dispatchAnalytics])
 
   // Refs
   const chatContainerRef = useRef<HTMLDivElement | null>(null)
@@ -194,16 +193,23 @@ const ChatV2Content = () => {
     }
 
     try {
-      const { tokenUsageAnalysis, stream } = await postCompletionStreamV3({
-        generationInfo,
-        messages: newMessages,
-        ragIndexId,
+      if (!streamController) {
+        throw new Error('streamController is not defined')
+      }
+
+      const { tokenUsageAnalysis, stream } = await postCompletionStreamV3(
         formData,
-        modelTemperature: parseFloat(modelTemperature),
-        courseId,
-        abortController: streamController,
-        saveConsent,
-      })
+        {
+          options: {
+            generationInfo,
+            chatMessages: newMessages,
+            modelTemperature,
+            saveConsent,
+          },
+          courseId,
+        },
+        streamController,
+      )
 
       if (!stream && !tokenUsageAnalysis) {
         console.error('getCompletionStream did not return a stream or token usage analysis')
@@ -263,7 +269,7 @@ const ChatV2Content = () => {
 
     const { usage, limit, model: defaultCourseModel, models: courseModels } = userStatus
 
-    let allowedModels: string[] = []
+    let allowedModels: ValidModelName[] = []
 
     if (course && courseModels) {
       allowedModels = courseModels
@@ -558,8 +564,8 @@ const ChatV2Content = () => {
       <SettingsModal
         open={settingsModalOpen}
         setOpen={setSettingsModalOpen}
-        modelTemperature={parseFloat(modelTemperature)}
-        setModelTemperature={(updatedTemperature) => setModelTemperature(String(updatedTemperature))}
+        modelTemperature={modelTemperature}
+        setModelTemperature={(updatedTemperature) => setModelTemperature(updatedTemperature)}
       />
 
       <DisclaimerModal disclaimer={disclaimerInfo} disclaimerStatus={disclaimerStatus} setDisclaimerStatus={setDisclaimerStatus} />
@@ -586,9 +592,9 @@ const LeftMenu = ({
   setSettingsModalOpen: React.Dispatch<React.SetStateAction<boolean>>
   setDisclaimerStatus: React.Dispatch<React.SetStateAction<boolean>>
   messages: ChatMessage[]
-  currentModel: string
-  setModel: (model: string) => void
-  availableModels: string[]
+  currentModel: ValidModelName
+  setModel: (model: ValidModelName) => void
+  availableModels: ValidModelName[]
 }) => {
   const { t } = useTranslation()
   const { courseId } = useParams()

src/client/components/ChatV2/ModelSelector.tsx

@@ -2,7 +2,7 @@ import React from 'react'
 import { useTranslation } from 'react-i18next'
 import { MenuItem, Typography, Tooltip, Menu } from '@mui/material'
 import { KeyboardArrowDown, SmartToy } from '@mui/icons-material'
-import { FREE_MODEL } from '../../../config'
+import { FREE_MODEL, ValidModelName } from '../../../config'
 import { OutlineButtonBlack } from './general/Buttons'
 
 const ModelSelector = ({
@@ -11,9 +11,9 @@ const ModelSelector = ({
   availableModels,
   isTokenLimitExceeded,
 }: {
-  currentModel: string
-  setModel: (model: string) => void
-  availableModels: string[]
+  currentModel: ValidModelName
+  setModel: (model: ValidModelName) => void
+  availableModels: ValidModelName[]
   isTokenLimitExceeded: boolean
 }) => {
   const { t } = useTranslation()
@@ -25,19 +25,14 @@ const ModelSelector = ({
     setAnchorEl(event.currentTarget)
   }
 
-  const handleSelect = (model: string) => {
+  const handleSelect = (model: ValidModelName) => {
     setModel(model)
     setAnchorEl(null)
   }
 
   return (
     <>
-      <OutlineButtonBlack
-        startIcon={<SmartToy />}
-        endIcon={<KeyboardArrowDown />}
-        onClick={handleClick}
-        data-testid="model-selector"
-      >
+      <OutlineButtonBlack startIcon={<SmartToy />} endIcon={<KeyboardArrowDown />} onClick={handleClick} data-testid="model-selector">
         {`${t('admin:model')}: ${validModel}`}
       </OutlineButtonBlack>
       <Menu
@@ -56,8 +51,10 @@ const ModelSelector = ({
           <MenuItem
             key={model}
             value={model}
-            onClick={() => handleSelect(model)} disabled={isTokenLimitExceeded && model !== FREE_MODEL}
-            data-testid={`${model}-option`}>
+            onClick={() => handleSelect(model)}
+            disabled={isTokenLimitExceeded && model !== FREE_MODEL}
+            data-testid={`${model}-option`}
+          >
             <Typography>
               {model}
               {model === FREE_MODEL && (

src/client/components/ChatV2/api.ts

@@ -1,4 +1,4 @@
-import type { ChatMessage, MessageGenerationInfo } from '../../../shared/chat'
+import type { PostStreamSchemaV3Type } from '../../../shared/chat'
 import { postAbortableStream } from '../../util/apiClient'
 import type { ChatToolOutput } from '../../../shared/tools'
 import { useGetQuery } from '../../hooks/apiHooks'
@@ -12,45 +12,8 @@ export const useToolResults = (toolCallId: string) => {
   })
 }
 
-interface PostCompletionStreamProps {
-  generationInfo: MessageGenerationInfo
-  courseId?: string
-  messages: ChatMessage[]
-  formData: FormData
-  ragIndexId?: number
-  userConsent?: boolean
-  modelTemperature: number
-  prevResponseId?: string
-  abortController?: AbortController
-  saveConsent: boolean
-}
-export const postCompletionStreamV3 = async ({
-  generationInfo,
-  courseId,
-  messages,
-  formData,
-  ragIndexId,
-  userConsent,
-  modelTemperature,
-  prevResponseId,
-  abortController,
-  saveConsent,
-}: PostCompletionStreamProps) => {
-  const data = {
-    courseId,
-    options: {
-      chatMessages: messages,
-      systemMessage: generationInfo.promptInfo.systemMessage,
-      model: generationInfo.model,
-      ragIndexId,
-      userConsent,
-      modelTemperature,
-      saveConsent,
-      prevResponseId,
-    },
-  }
-
-  formData.set('data', JSON.stringify(data))
+export const postCompletionStreamV3 = async (formData: FormData, input: PostStreamSchemaV3Type, abortController: AbortController) => {
+  formData.set('data', JSON.stringify(input))
 
   return postAbortableStream('/ai/v3/stream', formData, abortController)
 }

src/client/components/Courses/Course/EditCourseForm.tsx

@@ -7,7 +7,7 @@ import { enqueueSnackbar } from 'notistack'
 
 import { Course, SetState, User } from '../../../types'
 import { useEditCourseMutation } from '../../../hooks/useCourseMutation'
-import { validModels } from '../../../../config'
+import { ValidModelName, validModels } from '../../../../config'
 
 const EditCourseForm = forwardRef(({ course, setOpen, user }: { course: Course; setOpen: SetState<boolean>; user: User }, ref) => {
   const { t } = useTranslation()
@@ -64,7 +64,7 @@ const EditCourseForm = forwardRef(({ course, setOpen, user }: { course: Course;
               {t('admin:model')}
             </Typography>
             <Typography mb={1}>{t('admin:modelInfo')}</Typography>
-            <Select sx={{ m: 1, width: '300px' }} value={model} onChange={(e) => setModel(e.target.value)}>
+            <Select sx={{ m: 1, width: '300px' }} value={model} onChange={(e) => setModel(e.target.value as ValidModelName)}>
               {validModels.map(({ name: modelName }) => (
                 <MenuItem key={modelName} value={modelName}>
                   {modelName}

src/client/types.ts

@@ -1,3 +1,4 @@
+import { ValidModelName } from '../config'
 import type { ChatMessage } from '../shared/chat'
 import type { UserPreferences } from '../shared/user'
 
@@ -71,7 +72,7 @@ export interface ChatInstance {
   id: string
   name: Locales
   description: string
-  model: string
+  model: ValidModelName
   usageLimit: number
   resetCron?: string
   courseId?: string
@@ -83,7 +84,7 @@ export interface ChatInstance {
 export interface AccessGroup {
   id: string
   iamGroup: string
-  model: string
+  model: ValidModelName
   usageLimit: number | null
   resetCron: string | null
 }
@@ -139,8 +140,8 @@ export type ChatInstanceUsage = {
 }
 
 export type UserStatus = {
-  model: string
-  models: string[]
+  model: ValidModelName
+  models: ValidModelName[]
   usage: number
   limit: number
   isTike: boolean

src/config.ts

@@ -1,3 +1,5 @@
+import z from 'zod/v4'
+
 export const inDevelopment = process.env.NODE_ENV === 'development'
 
 export const inStaging = process.env.STAGING === 'true'
@@ -12,8 +14,6 @@ export const PUBLIC_URL = process.env.PUBLIC_URL || ''
 
 export const DEFAULT_TOKEN_LIMIT = Number(process.env.DEFAULT_TOKEN_LIMIT) || 150_000
 
-export const FREE_MODEL = process.env.FREE_MODEL || 'gpt-4o-mini' // as it was decided in 23th Sept 2024 meeting
-export const DEFAULT_MODEL = process.env.DEFAUL_MODEL || 'gpt-4o-mini'
 export const DEFAUL_CONTEXT_LIMIT = Number(process.env.DEFAUL_CONTEXT_LIMIT) || 4_096
 
 export const DEFAULT_RESET_CRON = process.env.DEFAULT_RESET_CRON || '0 0 1 */3 *'
@@ -41,7 +41,15 @@ export const validModels = [
     name: 'mock',
     context: 128_000,
   },
-]
+] as const
+
+export const ValidModelNameSchema = z.union(validModels.map((model) => z.literal(model.name)))
+
+export type ValidModelName = z.infer<typeof ValidModelNameSchema>
+
+export const DEFAULT_MODEL = ValidModelNameSchema.parse(process.env.DEFAULT_MODEL || 'gpt-4o-mini')
+
+export const FREE_MODEL = ValidModelNameSchema.parse(process.env.FREE_MODEL || 'gpt-4o-mini') // as it was decided in 23th Sept 2024 meeting
 
 export const DEFAULT_MODEL_ON_ENABLE = 'gpt-5'
 

src/server/db/models/prompt.ts

@@ -1,7 +1,8 @@
-import { type CreationOptional, DataTypes, type InferAttributes, type InferCreationAttributes, Model } from 'sequelize'
+import { type CreationOptional, DataTypes, type InferAttributes, type InferCreationAttributes, Model, NonAttribute } from 'sequelize'
 
 import type { CustomMessage } from '../../types'
 import { sequelize } from '../connection'
+import type RagIndex from './ragIndex'
 
 export const PromptTypeValues = ['CHAT_INSTANCE', 'PERSONAL'] as const
 export type PromptType = (typeof PromptTypeValues)[number]
@@ -26,6 +27,8 @@ class Prompt extends Model<InferAttributes<Prompt>, InferCreationAttributes<Prom
   declare hidden: CreationOptional<boolean>
 
   declare mandatory: CreationOptional<boolean>
+
+  declare ragIndex?: NonAttribute<RagIndex>
 }
 
 Prompt.init(