diff --git a/nx-dev/nx-dev/pages/api/query-ai-handler.ts b/nx-dev/nx-dev/pages/api/query-ai-handler.ts index 197158bdd94b50..c799fd2d40402a 100644 --- a/nx-dev/nx-dev/pages/api/query-ai-handler.ts +++ b/nx-dev/nx-dev/pages/api/query-ai-handler.ts @@ -8,14 +8,15 @@ import { PROMPT, PageSection, appendToStream, - checkEnvVariables, + getSupabaseClient, formatMarkdownSources, getLastAssistantMessageContent, + getOpenAI, getUserQuery, initializeChat, // moderateContent, } from '@nx/nx-dev/util-ai'; -import { SupabaseClient, createClient } from '@supabase/supabase-js'; +import { SupabaseClient } from '@supabase/supabase-js'; import OpenAI from 'openai'; import { OpenAIStream, StreamingTextResponse } from 'ai'; import GPT3Tokenizer from 'gpt3-tokenizer'; @@ -24,6 +25,10 @@ import { Stream } from 'openai/streaming'; const supabaseUrl = process.env['NX_NEXT_PUBLIC_SUPABASE_URL']; const supabaseServiceKey = process.env['NX_SUPABASE_SERVICE_ROLE_KEY_ACTUAL']; const openAiKey = process.env['NX_OPENAI_KEY']; +const tokenCountLimit = + parseInt(process.env['NX_TOKEN_COUNT_LIMIT'] ?? '2500') > 0 + ? parseInt(process.env['NX_TOKEN_COUNT_LIMIT'] ?? '2500') + : 2500; export const config = { runtime: 'edge', @@ -31,18 +36,12 @@ export const config = { export default async function handler(request: NextRequest) { try { - checkEnvVariables(openAiKey, supabaseUrl, supabaseServiceKey); - const openai = new OpenAI({ - apiKey: openAiKey, // This is also the default, can be omitted - }); + const openai = getOpenAI(openAiKey); + const supabaseClient: SupabaseClient = + getSupabaseClient(supabaseUrl, supabaseServiceKey); const { messages } = (await request.json()) as { messages: ChatItem[] }; - const supabaseClient: SupabaseClient = createClient( - supabaseUrl as string, - supabaseServiceKey as string - ); - const query: string | null = getUserQuery(messages); const sanitizedQuery = query.trim(); @@ -102,7 +101,7 @@ export default async function handler(request: NextRequest) { const encoded = tokenizer.encode(content); tokenCount += encoded.text.length; - if (tokenCount >= 2500) { + if (tokenCount >= tokenCountLimit) { break; } diff --git a/nx-dev/util-ai/package.json b/nx-dev/util-ai/package.json index 74410b250a84eb..dd624378852c91 100644 --- a/nx-dev/util-ai/package.json +++ b/nx-dev/util-ai/package.json @@ -2,6 +2,7 @@ "name": "@nx/nx-dev/util-ai", "version": "0.0.1", "dependencies": { + "@supabase/supabase-js": "^2.26.0", "tslib": "^2.3.0", "openai": "~4.3.1" }, diff --git a/nx-dev/util-ai/src/lib/chat-utils.ts b/nx-dev/util-ai/src/lib/chat-utils.ts index 21e4ca0d4ad503..1d1c93e1c127c8 100644 --- a/nx-dev/util-ai/src/lib/chat-utils.ts +++ b/nx-dev/util-ai/src/lib/chat-utils.ts @@ -125,21 +125,14 @@ export async function appendToStream( originalStream: ReadableStream, appendContent: string ): Promise> { - let controller: ReadableStreamDefaultController; - return new ReadableStream({ - async start(ctrl) { - controller = ctrl; - const reader = originalStream.getReader(); - - let result; - while (!(result = await reader.read()).done) { - controller.enqueue(result.value); - } - - controller.enqueue(new TextEncoder().encode(appendContent)); - controller.close(); + const appendText = new TransformStream({ + flush(ctrl) { + ctrl.enqueue(new TextEncoder().encode(appendContent)); + ctrl.terminate(); }, }); + + return originalStream.pipeThrough(appendText); } export function getLastAssistantIndex(messages: ChatItem[]): number { diff --git a/nx-dev/util-ai/src/lib/utils.ts b/nx-dev/util-ai/src/lib/utils.ts index 889d3faf575110..b86da7eab8d59f 100644 --- a/nx-dev/util-ai/src/lib/utils.ts +++ b/nx-dev/util-ai/src/lib/utils.ts @@ -1,18 +1,10 @@ +import OpenAI from 'openai'; +import { SupabaseClient, createClient } from '@supabase/supabase-js'; + export function checkEnvVariables( - openAiKey?: string, supabaseUrl?: string, supabaseServiceKey?: string ) { - if (!openAiKey) { - throw new CustomError( - 'application_error', - 'Missing environment variable NX_OPENAI_KEY', - { - missing_key: true, - } - ); - } - if (!supabaseUrl) { throw new CustomError( 'application_error', @@ -29,6 +21,37 @@ export function checkEnvVariables( } } +let openai: OpenAI; +let supabaseClient: SupabaseClient; + +export function getOpenAI(openAiKey?: string): OpenAI { + if (openai) return openai; + if (!openAiKey) { + throw new CustomError( + 'application_error', + 'Missing environment variable NX_OPENAI_KEY', + { + missing_key: true, + } + ); + } + openai = new OpenAI({ apiKey: openAiKey }); + return openai; +} + +export function getSupabaseClient( + supabaseUrl?: string, + supabaseServiceKey?: string +): SupabaseClient { + if (supabaseClient) return supabaseClient; + checkEnvVariables(supabaseUrl, supabaseServiceKey); + supabaseClient = createClient( + supabaseUrl as string, + supabaseServiceKey as string + ); + return supabaseClient; +} + export class CustomError extends Error { public type: string; public data: Record;