Changes to move away from edge endpoint

USEPA · Oct 10, 2024 · 9770c8f · 9770c8f
1 parent c6c912c
commit 9770c8f
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 99 deletions.
diff --git a/pages/api/chat.ts b/pages/api/chat.ts
@@ -2,43 +2,26 @@ import { DEFAULT_SYSTEM_PROMPT, DEFAULT_TEMPERATURE } from '@/utils/app/const';
 import { OpenAIError, OpenAIStream } from '@/utils/server';
 import { ChatBody, Message } from '@/types/chat';
 
-// @ts-expect-error
-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 { NextApiRequest, NextApiResponse } from 'next';
+import { Tiktoken } from '@dqbd/tiktoken';
 
- export const config = {
-   runtime: 'edge',
- };
-
-const handler = async (req: Request): Promise<Response> => {
+const handler = async (req: NextApiRequest, res: NextApiResponse<any>) => {
   try {
-    const { model, messages, key, prompt, temperature } = (await req.json()) as ChatBody;
-
-    await init((imports) => WebAssembly.instantiate(wasm, imports));
+    const { model, messages, key, prompt, temperature } = req.body as ChatBody;
     const encoding = new Tiktoken(
       tiktokenModel.bpe_ranks,
       tiktokenModel.special_tokens,
       tiktokenModel.pat_str,
     );
-
-    let promptToSend = prompt;
-    if (!promptToSend) {
-      promptToSend = DEFAULT_SYSTEM_PROMPT;
-    }
-
-    let temperatureToUse = temperature;
-    if (temperatureToUse == null) {
-      temperatureToUse = DEFAULT_TEMPERATURE;
-    }
+    let promptToSend = prompt || DEFAULT_SYSTEM_PROMPT;
+    let temperatureToUse = temperature ?? DEFAULT_TEMPERATURE;
 
     const prompt_tokens = encoding.encode(promptToSend);
-
     let tokenCount = prompt_tokens.length;
     let messagesToSend: Message[] = [];
 
+    // Reverse loop through the messages to add them until the token limit is reached
     for (let i = messages.length - 1; i >= 0; i--) {
       const message = messages[i];
       const tokens = encoding.encode(message.content);
@@ -49,26 +32,54 @@ const handler = async (req: Request): Promise<Response> => {
       tokenCount += tokens.length;
       messagesToSend = [message, ...messagesToSend];
     }
-    // console.log("!!!HEADERS!!!");
-    // console.log(req.headers);
-    var principalName:string|null = req.headers.get("x-ms-client-principal-name");
-    var bearer:string|null =req.headers.get("x-ms-token-aad-access-token")? req.headers.get("x-ms-token-aad-access-token") : req.headers.get("x-ms-client-principal");
-    var bearerAuth: string|null = req.headers.get("x-ms-client-principal-id");
-    const userName = req.headers.get("x-ms-client-principal-name")
-    // console.log("principalName:" + principalName);
-    // console.log("bearer:" + bearer);
+
+    const principalName: string = req.headers['x-ms-client-principal-name']?.toString() || "";
+    const bearer: string = req.headers['x-ms-token-aad-access-token']?.toString() || req.headers['x-ms-client-principal']?.toString() || "";
+    const bearerAuth: string = req.headers['x-ms-client-principal-id']?.toString() || "";
+    const userName: string = req.headers['x-ms-client-principal-name']?.toString() || "";
+
     encoding.free();
 
+    const stream = await OpenAIStream(
+      model,
+      promptToSend,
+      temperatureToUse,
+      key,
+      messagesToSend,
+      principalName,
+      bearer,
+      bearerAuth,
+      userName
+    );
+
+    res.setHeader('Content-Type', 'text/event-stream');
+    res.setHeader('Cache-Control', 'no-cache');
+    res.setHeader('Connection', 'keep-alive');
+
+    const reader = stream.getReader();
+    const decoder = new TextDecoder();
+
+    const processStream = async () => {
+      let done = false;
+      while (!done) {
+        const { value, done: readerDone } = await reader.read();
+        done = readerDone;
+        if (value) {
+          const chunk = decoder.decode(value, { stream: !done });
+          res.write(chunk);
+        }
+      }
+      res.end();
+    };
 
-    const stream = await OpenAIStream(model, promptToSend, temperatureToUse, key, messagesToSend, principalName, bearer, bearerAuth, userName );
+    await processStream();
 
-    return new Response(stream);
   } catch (error) {
     console.error(error);
     if (error instanceof OpenAIError) {
-      return new Response('Error', { status: 500, statusText: error.message });
+      res.status(500).send(error.message);
     } else {
-      return new Response('Error', { status: 500 });
+      res.status(500).send('Internal Server Error');
     }
   }
 };

diff --git a/pages/api/health-check.tsx b/pages/api/health-check.tsx
@@ -1,55 +1,39 @@
-import { NextRequest, NextResponse } from 'next/server';
-import handler from './chat';  // Adjust this path based on the actual location of 'chat'
+import { NextApiRequest, NextApiResponse } from 'next';
+import handler from './chat';
 
-// Declare the function as an edge runtime function
-export const config = {
-  runtime: 'edge', // Ensure this is executed in Edge runtime
-};
-
-export default async function healthCheck(req: NextRequest) {
+export default async function healthCheck(req: NextApiRequest, res: NextApiResponse) {
   try {
-    const mockReq: Request = new Request("http://example.org", {
-      method: 'POST',
-      headers: {
-        'content-type': 'application/json',
-        'x-ms-client-principal-name': req.headers?.get("x-ms-client-principal-name") || '',
-        'x-ms-token-aad-access-token': req.headers?.get("x-ms-token-aad-access-token") || '',
-        'x-ms-client-principal-id': req.headers?.get("x-ms-client-principal-id") || ''
+    req.method = 'POST';
+    req.headers['content-type'] = 'application/json';
+    req.headers['x-ms-client-principal-name'] = req.headers['x-ms-client-principal-name'] || '';
+    req.headers['x-ms-token-aad-access-token'] = req.headers['x-ms-token-aad-access-token'] || '';
+    req.headers['x-ms-client-principal-id'] = req.headers['x-ms-client-principal-id'] || '';
+
+    req.body = {
+      model: {
+        id: 'gpt-35-turbo',
+        name: 'GPT-3.5',
+        maxLength: 12000,
+        tokenLimit: 4000,
       },
-      body: JSON.stringify({
-        model: {
-          id: 'gpt-35-turbo',
-          name: 'GPT-3.5',
-          maxLength: 12000,
-          tokenLimit: 4000,
+      messages: [
+        {
+          role: 'user',
+          content: 'test',
         },
-        messages: [
-          {
-            role: 'user',
-            content: 'test',
-          },
-        ],
-        key: '',
-        prompt: "You are Gov Chat an AI Assistant using Azure OpenAI. Follow the user's instructions carefully. Respond using markdown.",
-        temperature: 0.5,
-      }),
-    });
-
-    // Simulate the handler with the mock request
-    const mockRes:Response = await handler(mockReq);
-    if (mockRes.status != 200) {
-      throw Error("Status not 200")
-    }
-
-    // Return a successful health check
-    return NextResponse.json({ message: 'Health check passed successfully' }, { status: 200 });
+      ],
+      key: '',
+      prompt: "You are Gov Chat an AI Assistant using Azure OpenAI. Follow the user's instructions carefully. Respond using markdown.",
+      temperature: 0.5,
+    };
+    await handler(req, res);
   } catch (error: unknown) {
     if (error instanceof Error) {
       console.error('Health check failed:', error.message);
-      return NextResponse.json({ message: 'Health check failed', error: error.message }, { status: 500 });
+      return res.status(500).send(error.message );
     } else {
       // Handle cases where a non-Error object is thrown
-      return NextResponse.json({ message: 'Unknown error occurred' }, { status: 500 });
+      return res.status(500).send('Unknown error occurred');
     }
   }
 }
diff --git a/pages/api/models.ts b/pages/api/models.ts
@@ -1,14 +1,12 @@
 import { OPENAI_API_HOST, OPENAI_API_TYPE, OPENAI_API_VERSION, OPENAI_ORGANIZATION, AZURE_APIM } from '@/utils/app/const';
 import { OpenAIModel, OpenAIModelID, OpenAIModels } from '@/types/openai';
 import { getAuthToken } from '@/utils/lib/azure';
+import type { NextApiRequest, NextApiResponse } from 'next';
 
-export const config = {
-  runtime: 'edge',
-};
 
-const handler = async (req: Request): Promise<Response> => {
+const handler = async (req: NextApiRequest, res: NextApiResponse) => {
   try {
-    const { key } = (await req.json()) as {
+    const { key } = req.body as {
       key: string;
     };
 
@@ -18,9 +16,9 @@ const handler = async (req: Request): Promise<Response> => {
     }
 
     let token = await getAuthToken();
-
+    // TODO: Seems to currently not be working. Also doesn't seem to be used right now.
     const response = await fetch(url, {
-      method: 'post',
+      method: 'POST',
       headers: {
         'Content-Type': 'application/json',
         ...(OPENAI_API_TYPE === 'openai' && {
@@ -42,10 +40,8 @@ const handler = async (req: Request): Promise<Response> => {
     });
 
     if (response.status === 401) {
-      return new Response(response.body, {
-        status: 500,
-        headers: response.headers,
-      });
+      res.status(500).json({ error: 'Unauthorized' });
+      return;
     } else if (response.status !== 200) {
       console.error(
         `OpenAI API returned an error ${response.status
@@ -69,11 +65,11 @@ const handler = async (req: Request): Promise<Response> => {
         }
       })
       .filter(Boolean);
-      console.log("loaded models");
-    return new Response(JSON.stringify(models), { status: 200 });
+    console.log("loaded models");
+    res.status(200).json(models);
   } catch (error) {
     console.error(error);
-    return new Response('Error', { status: 500 });
+    res.status(500).send('Error');
   }
 };
 

diff --git a/pages/home/home.tsx b/pages/home/home.tsx
@@ -86,12 +86,6 @@ const Home = ({
     queryFn: ({ signal }) => {
       if (!apiKey && !serverSideApiKeyIsSet) return null;
       return {};
-      return getModels(
-        {
-          key: apiKey,
-        },
-        signal,
-      );
     },
     enabled: true,
     refetchOnMount: false