fix(nx-dev): address comments on PR

nrwl · Aug 30, 2023 · a285673 · a285673
1 parent dcdf1a7
commit a285673
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 36 deletions.
diff --git 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,25 +25,23 @@ 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',
 };
 
 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<any, 'public', any> =
+      getSupabaseClient(supabaseUrl, supabaseServiceKey);
 
     const { messages } = (await request.json()) as { messages: ChatItem[] };
 
-    const supabaseClient: SupabaseClient<any, 'public', any> = 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
@@ -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
@@ -125,21 +125,14 @@ export async function appendToStream(
   originalStream: ReadableStream<Uint8Array>,
   appendContent: string
 ): Promise<ReadableStream<Uint8Array>> {
-  let controller: ReadableStreamDefaultController<Uint8Array>;
-  return new ReadableStream<Uint8Array>({
-    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
@@ -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<any, 'public', any>;
+
+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<any, 'public', any> {
+  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<string, any>;