Fix RAG chat temperature handling

ui/__tests__/rag-chat-options.test.ts

@@ -0,0 +1,24 @@
+import { DEFAULT_TEMPERATURE } from '@/utils/app/const';
+import { resolveRagChatTemperature } from '@/utils/server/rag-chat-options';
+
+import { describe, expect, it } from 'vitest';
+
+describe('resolveRagChatTemperature', () => {
+  it('keeps an explicit caller-provided temperature', () => {
+    expect(resolveRagChatTemperature(0.7)).toBe(0.7);
+  });
+
+  it('falls back to the default temperature when omitted', () => {
+    expect(resolveRagChatTemperature(undefined)).toBe(DEFAULT_TEMPERATURE);
+    expect(resolveRagChatTemperature(null)).toBe(DEFAULT_TEMPERATURE);
+  });
+
+  it('falls back to the default temperature when the value is invalid', () => {
+    expect(resolveRagChatTemperature(Number.NaN)).toBe(DEFAULT_TEMPERATURE);
+    expect(resolveRagChatTemperature(Number.POSITIVE_INFINITY)).toBe(
+      DEFAULT_TEMPERATURE,
+    );
+    expect(resolveRagChatTemperature(-0.1)).toBe(DEFAULT_TEMPERATURE);
+    expect(resolveRagChatTemperature(2.1)).toBe(DEFAULT_TEMPERATURE);
+  });
+});

ui/pages/api/rag-chat.ts

@@ -1,6 +1,6 @@
-import { DEFAULT_SYSTEM_PROMPT, DEFAULT_TEMPERATURE } from '@/utils/app/const';
+import { DEFAULT_SYSTEM_PROMPT } from '@/utils/app/const';
 import { OpenAIError, OpenAIStream } from '@/utils/server';
-import { codeBlock, oneLine } from 'common-tags'
+import { resolveRagChatTemperature } from '@/utils/server/rag-chat-options';
 
 import { ChatBody, Message } from '@/types/chat';
 
@@ -9,6 +9,7 @@ import wasm from '../../node_modules/@dqbd/tiktoken/lite/tiktoken_bg.wasm?module
 
 import tiktokenModel from '@dqbd/tiktoken/encoders/cl100k_base.json';
 import { Tiktoken, init } from '@dqbd/tiktoken/lite/init';
+import { codeBlock, oneLine } from 'common-tags';
 
 export const config = {
   runtime: 'edge',
@@ -17,38 +18,36 @@ export const config = {
 // Function to fetch and format documents
 async function fetchAndFormatDocuments(lastMessageContent: string) {
   try {
-    console.log("fetching documents")
+    console.log('fetching documents');
     const response = await fetch('http://localhost:3000/api/fetch-documents', {
       method: 'POST',
       headers: { 'Content-Type': 'application/json' },
       body: JSON.stringify({ input: lastMessageContent }),
     });
-    
+
     if (!response.ok) {
       throw new Error(`Error fetching documents: ${response.statusText}`);
     }
 
     const data = await response.json();
-    const result = data.metadatas[0].map((metadata: any, index: number) => {
-      return `Source ${index + 1}) Title: ${metadata.title}, Page: ${metadata.page}, Content: ${data.documents[0][index]}\n`;
-    }).join('');
+    const result = data.metadatas[0]
+      .map((metadata: any, index: number) => {
+        return `Source ${index + 1}) Title: ${metadata.title}, Page: ${
+          metadata.page
+        }, Content: ${data.documents[0][index]}\n`;
+      })
+      .join('');
 
     console.log(result);
 
     return result;
-
   } catch (error) {
     console.error('Error fetching and formatting documents:', error);
     throw error; // You may want to throw a more specific error object here
   }
 }
 
-
-
-
-
 const handler = async (req: Request): Promise<Response> => {
-
   try {
     const { model, messages, key, prompt, temperature } =
       (await req.json()) as ChatBody;
@@ -85,34 +84,31 @@ const handler = async (req: Request): Promise<Response> => {
 
     const lastMessage = messages[messages.length - 1];
 
-    const relevantDocuments = await fetchAndFormatDocuments(lastMessage.content);
-
-    let temperatureToUse = temperature;
-    if (temperatureToUse == null) {
-      temperatureToUse = DEFAULT_TEMPERATURE;
-    }
+    const relevantDocuments = await fetchAndFormatDocuments(
+      lastMessage.content,
+    );
+
+    const temperatureToUse = resolveRagChatTemperature(temperature);
 
     const prompt_tokens = encoding.encode(promptToSend);
 
     let tokenCount = prompt_tokens.length;
     let messagesToSend: Message[] = [];
 
-
     encoding.free();
 
     console.log(model, promptToSend, temperatureToUse, key, messagesToSend);
 
-
-  messagesToSend = [
+    messagesToSend = [
       {
-        role: "user",
+        role: 'user',
         content: codeBlock`
           Here is the relevant documentation:
           ${relevantDocuments}
         `,
       },
       {
-        role: "user",
+        role: 'user',
         content: codeBlock`
           ${oneLine`
             Answer my next question using only the above documentation.
@@ -135,19 +131,18 @@ const handler = async (req: Request): Promise<Response> => {
         `,
       },
       {
-        role: "user",
+        role: 'user',
         content: codeBlock`
           Here is my question:
           ${oneLine`${lastMessage.content}`}
       `,
       },
-    ]
-
+    ];
 
     const stream = await OpenAIStream(
       model,
       promptToSend,
-      0,
+      temperatureToUse,
       key,
       messagesToSend,
     );

ui/utils/server/rag-chat-options.ts

@@ -0,0 +1,17 @@
+import { DEFAULT_TEMPERATURE } from '@/utils/app/const';
+
+const MIN_TEMPERATURE = 0;
+const MAX_TEMPERATURE = 2;
+
+export const resolveRagChatTemperature = (temperature?: number | null) => {
+  if (
+    temperature == null ||
+    !Number.isFinite(temperature) ||
+    temperature < MIN_TEMPERATURE ||
+    temperature > MAX_TEMPERATURE
+  ) {
+    return DEFAULT_TEMPERATURE;
+  }
+
+  return temperature;
+};