diff --git a/pages/api/chat.ts b/pages/api/chat.ts index f173359759..f1ed3ad305 100644 --- 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 => { +const handler = async (req: NextApiRequest, res: NextApiResponse) => { 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 => { 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 index bd45388a1b..fbe5dc3824 100644 --- 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 index 07f580dfd3..ac906ad517 100644 --- 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 => { +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 => { } 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 => { }); 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 => { } }) .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 index 2d10bd754f..2b7a695adb 100644 --- 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