Improve RAG context budgeting

ui/__tests__/rag-context.test.ts

@@ -0,0 +1,92 @@
+import {
+  DEFAULT_RAG_CONTEXT_CHAR_LIMIT,
+  formatRagSource,
+  parsePositiveInteger,
+  parseRagContextCharLimit,
+  prepareRagContext,
+} from '@/utils/server/rag-context';
+
+import { describe, expect, it } from 'vitest';
+
+describe('RAG context preparation', () => {
+  it('formats retrieved sources with metadata and distance', () => {
+    expect(
+      formatRagSource(
+        '  The method uses a controlled cohort.  ',
+        { title: 'Paper A', page: 4 },
+        2,
+        0.12891,
+      ),
+    ).toBe(
+      'Source 2) Title: Paper A, Page: 4, Distance: 0.1289, Content: The method uses a controlled cohort.\n',
+    );
+  });
+
+  it('deduplicates repeated chunks before building the prompt context', () => {
+    const result = prepareRagContext({
+      documents: [
+        ['Same abstract text', 'Same   abstract text', 'Different result'],
+      ],
+      metadatas: [
+        [
+          { title: 'Paper', page: 1 },
+          { title: 'Paper', page: 1 },
+          { title: 'Paper', page: 2 },
+        ],
+      ],
+      distances: [[0.1, 0.2, 0.3]],
+    });
+
+    expect(result.sourceCount).toBe(2);
+    expect(result.omittedSourceCount).toBe(1);
+    expect(result.context).toContain('Source 1) Title: Paper, Page: 1');
+    expect(result.context).toContain('Source 2) Title: Paper, Page: 2');
+  });
+
+  it('keeps the generated context under the configured character budget', () => {
+    const result = prepareRagContext(
+      {
+        documents: [
+          [
+            'A'.repeat(DEFAULT_RAG_CONTEXT_CHAR_LIMIT),
+            'This second source should be omitted after the context is full.',
+          ],
+        ],
+        metadatas: [
+          [
+            { title: 'Long Paper', page: 10 },
+            { title: 'Other Paper', page: 11 },
+          ],
+        ],
+      },
+      260,
+    );
+
+    expect(result.context.length).toBeLessThanOrEqual(260);
+    expect(result.sourceCount).toBe(1);
+    expect(result.omittedSourceCount).toBe(1);
+    expect(result.context).toContain('...');
+  });
+
+  it('returns an explicit empty-context message when no documents match', () => {
+    const result = prepareRagContext({ documents: [[]], metadatas: [[]] });
+
+    expect(result.sourceCount).toBe(0);
+    expect(result.context).toContain('No matching uploaded-document context');
+  });
+
+  it('parses bounded positive integer configuration values', () => {
+    expect(parsePositiveInteger('12', 8, 20)).toBe(12);
+    expect(parsePositiveInteger(100, 8, 20)).toBe(20);
+    expect(parsePositiveInteger('bad', 8, 20)).toBe(8);
+    expect(parsePositiveInteger(0, 8, 20)).toBe(8);
+  });
+
+  it('parses the configurable RAG context character budget', () => {
+    expect(parseRagContextCharLimit('16000')).toBe(16000);
+    expect(parseRagContextCharLimit('bad')).toBe(
+      DEFAULT_RAG_CONTEXT_CHAR_LIMIT,
+    );
+    expect(parseRagContextCharLimit(100000)).toBe(50000);
+  });
+});

ui/pages/api/fetch-documents.ts

@@ -1,25 +1,64 @@
-import type { NextApiRequest, NextApiResponse } from "next";
-import { ChromaClient, TransformersEmbeddingFunction } from "chromadb";
+import type { NextApiRequest, NextApiResponse } from 'next';
 
-export default async function handler(req: NextApiRequest, res: NextApiResponse) {
+import {
+  DEFAULT_RAG_RETRIEVAL_RESULTS,
+  parsePositiveInteger,
+  parseRagContextCharLimit,
+  prepareRagContext,
+} from '@/utils/server/rag-context';
+
+import { ChromaClient, TransformersEmbeddingFunction } from 'chromadb';
+
+export default async function handler(
+  req: NextApiRequest,
+  res: NextApiResponse,
+) {
   try {
+    if (req.method !== 'POST') {
+      return res.status(405).end();
+    }
+
     const client = new ChromaClient({
-      path: "http://chroma-server:8000",
+      path: process.env.CHROMA_PATH || 'http://chroma-server:8000',
     });
 
     const query = req.body.input;
 
+    if (typeof query !== 'string' || query.trim().length === 0) {
+      return res
+        .status(400)
+        .json({ error: 'A non-empty input query is required' });
+    }
+
+    const nResults = parsePositiveInteger(
+      req.body.nResults,
+      DEFAULT_RAG_RETRIEVAL_RESULTS,
+      20,
+    );
+    const contextCharLimit = parseRagContextCharLimit(
+      req.body.contextCharLimit ?? process.env.RAG_CONTEXT_CHAR_LIMIT,
+    );
+
     const embedder = new TransformersEmbeddingFunction();
 
-    const collection = await client.getOrCreateCollection({ name: "default-collection", embeddingFunction: embedder });
+    const collection = await client.getOrCreateCollection({
+      name: 'default-collection',
+      embeddingFunction: embedder,
+    });
 
-  // query the collection
-  const results = await collection.query({
-      nResults: 4, 
-      queryTexts: [query]
-  }) 
+    // query the collection
+    const results = await collection.query({
+      nResults,
+      queryTexts: [query.trim()],
+    });
 
-    res.status(200).json(results);
+    const preparedContext = prepareRagContext(results, contextCharLimit);
+
+    res.status(200).json({
+      ...results,
+      _prepared: preparedContext,
+      ...preparedContext,
+    });
   } catch (error) {
     if (error instanceof Error) {
       console.error('Error message:', error.message);
@@ -29,4 +68,4 @@ export default async function handler(req: NextApiRequest, res: NextApiResponse)
     }
     res.status(500).json({ error: 'An unexpected error occurred :(' });
   }
-}
+}

ui/pages/api/rag-chat.ts

@@ -1,6 +1,6 @@
 import { DEFAULT_SYSTEM_PROMPT, DEFAULT_TEMPERATURE } from '@/utils/app/const';
 import { OpenAIError, OpenAIStream } from '@/utils/server';
-import { codeBlock, oneLine } from 'common-tags'
+import { prepareRagContext } from '@/utils/server/rag-context';
 
 import { ChatBody, Message } from '@/types/chat';
 
@@ -9,46 +9,48 @@ 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',
 };
 
 // Function to fetch and format documents
-async function fetchAndFormatDocuments(lastMessageContent: string) {
+async function fetchAndFormatDocuments(
+  lastMessageContent: string,
+  requestOrigin: string,
+) {
   try {
-    console.log("fetching documents")
-    const response = await fetch('http://localhost:3000/api/fetch-documents', {
+    console.log('fetching documents');
+    const fetchDocumentsUrl = new URL('/api/fetch-documents', requestOrigin);
+    const response = await fetch(fetchDocumentsUrl, {
       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 =
+      typeof data._prepared?.context === 'string'
+        ? data._prepared.context
+        : typeof data.context === 'string'
+        ? data.context
+        : prepareRagContext(data).context;
 
     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
+    return 'No matching uploaded-document context was found for this question.';
   }
 }
 
-
-
-
-
 const handler = async (req: Request): Promise<Response> => {
-
   try {
     const { model, messages, key, prompt, temperature } =
       (await req.json()) as ChatBody;
@@ -85,8 +87,11 @@ const handler = async (req: Request): Promise<Response> => {
 
     const lastMessage = messages[messages.length - 1];
 
-    const relevantDocuments = await fetchAndFormatDocuments(lastMessage.content);
-
+    const relevantDocuments = await fetchAndFormatDocuments(
+      lastMessage.content,
+      new URL(req.url).origin,
+    );
+
     let temperatureToUse = temperature;
     if (temperatureToUse == null) {
       temperatureToUse = DEFAULT_TEMPERATURE;
@@ -97,22 +102,20 @@ const handler = async (req: Request): Promise<Response> => {
     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,14 +138,13 @@ 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,