fix: api-key manager cleanup and log error on llm call (#1077)

* fix: api-key manager cleanup and log error on llm call * log improved
stackblitz-labs · Jan 12, 2025 · fad4197 · fad4197
1 parent 3a298f1
commit fad4197
Show file tree

Hide file tree

Showing 6 changed files with 65 additions and 44 deletions.
diff --git a/app/components/chat/APIKeyManager.tsx b/app/components/chat/APIKeyManager.tsx
@@ -2,7 +2,6 @@ import React, { useState, useEffect, useCallback } from 'react';
 import { IconButton } from '~/components/ui/IconButton';
 import type { ProviderInfo } from '~/types/model';
 import Cookies from 'js-cookie';
-import { providerBaseUrlEnvKeys } from '~/utils/constants';
 
 interface APIKeyManagerProps {
   provider: ProviderInfo;
@@ -93,17 +92,15 @@ export const APIKeyManager: React.FC<APIKeyManagerProps> = ({ provider, apiKey,
           <span className="text-sm font-medium text-bolt-elements-textSecondary">{provider?.name} API Key:</span>
           {!isEditing && (
             <div className="flex items-center gap-2">
-              {isEnvKeySet ? (
+              {apiKey ? (
                 <>
                   <div className="i-ph:check-circle-fill text-green-500 w-4 h-4" />
-                  <span className="text-xs text-green-500">
-                    Set via {providerBaseUrlEnvKeys[provider.name].apiTokenKey} environment variable
-                  </span>
+                  <span className="text-xs text-green-500">Set via UI</span>
                 </>
-              ) : apiKey ? (
+              ) : isEnvKeySet ? (
                 <>
                   <div className="i-ph:check-circle-fill text-green-500 w-4 h-4" />
-                  <span className="text-xs text-green-500">Set via UI</span>
+                  <span className="text-xs text-green-500">Set via environment variable</span>
                 </>
               ) : (
                 <>
@@ -117,7 +114,7 @@ export const APIKeyManager: React.FC<APIKeyManagerProps> = ({ provider, apiKey,
       </div>
 
       <div className="flex items-center gap-2 shrink-0">
-        {isEditing && !isEnvKeySet ? (
+        {isEditing ? (
           <div className="flex items-center gap-2">
             <input
               type="password"
@@ -145,16 +142,16 @@ export const APIKeyManager: React.FC<APIKeyManagerProps> = ({ provider, apiKey,
           </div>
         ) : (
           <>
-            {!isEnvKeySet && (
+            {
               <IconButton
                 onClick={() => setIsEditing(true)}
                 title="Edit API Key"
                 className="bg-blue-500/10 hover:bg-blue-500/20 text-blue-500"
               >
                 <div className="i-ph:pencil-simple w-4 h-4" />
               </IconButton>
-            )}
-            {provider?.getApiKeyLink && !isEnvKeySet && (
+            }
+            {provider?.getApiKeyLink && !apiKey && (
               <IconButton
                 onClick={() => window.open(provider?.getApiKeyLink)}
                 title="Get API Key"

diff --git a/app/components/chat/Chat.client.tsx b/app/components/chat/Chat.client.tsx
@@ -137,35 +137,36 @@ export const ChatImpl = memo(
 
     const [apiKeys, setApiKeys] = useState<Record<string, string>>({});
 
-    const { messages, isLoading, input, handleInputChange, setInput, stop, append, setMessages, reload } = useChat({
-      api: '/api/chat',
-      body: {
-        apiKeys,
-        files,
-        promptId,
-        contextOptimization: contextOptimizationEnabled,
-      },
-      sendExtraMessageFields: true,
-      onError: (error) => {
-        logger.error('Request failed\n\n', error);
-        toast.error(
-          'There was an error processing your request: ' + (error.message ? error.message : 'No details were returned'),
-        );
-      },
-      onFinish: (message, response) => {
-        const usage = response.usage;
-
-        if (usage) {
-          console.log('Token usage:', usage);
-
-          // You can now use the usage data as needed
-        }
+    const { messages, isLoading, input, handleInputChange, setInput, stop, append, setMessages, reload, error } =
+      useChat({
+        api: '/api/chat',
+        body: {
+          apiKeys,
+          files,
+          promptId,
+          contextOptimization: contextOptimizationEnabled,
+        },
+        sendExtraMessageFields: true,
+        onError: (e) => {
+          logger.error('Request failed\n\n', e, error);
+          toast.error(
+            'There was an error processing your request: ' + (e.message ? e.message : 'No details were returned'),
+          );
+        },
+        onFinish: (message, response) => {
+          const usage = response.usage;
 
-        logger.debug('Finished streaming');
-      },
-      initialMessages,
-      initialInput: Cookies.get(PROMPT_COOKIE_KEY) || '',
-    });
+          if (usage) {
+            console.log('Token usage:', usage);
+
+            // You can now use the usage data as needed
+          }
+
+          logger.debug('Finished streaming');
+        },
+        initialMessages,
+        initialInput: Cookies.get(PROMPT_COOKIE_KEY) || '',
+      });
     useEffect(() => {
       const prompt = searchParams.get('prompt');
 
@@ -263,6 +264,10 @@ export const ChatImpl = memo(
        */
       await workbenchStore.saveAllFiles();
 
+      if (error != null) {
+        setMessages(messages.slice(0, -1));
+      }
+
       const fileModifications = workbenchStore.getFileModifcations();
 
       chatStore.setKey('aborted', false);

diff --git a/app/lib/.server/llm/stream-text.ts b/app/lib/.server/llm/stream-text.ts
@@ -226,7 +226,7 @@ export async function streamText(props: {
 
   logger.info(`Sending llm call to ${provider.name} with model ${modelDetails.name}`);
 
-  return _streamText({
+  return await _streamText({
     model: provider.getModelInstance({
       model: currentModel,
       serverEnv,

diff --git a/app/routes/api.chat.ts b/app/routes/api.chat.ts
@@ -122,6 +122,8 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
         return;
       },
     };
+    const totalMessageContent = messages.reduce((acc, message) => acc + message.content, '');
+    logger.debug(`Total message length: ${totalMessageContent.split(' ').length}, words`);
 
     const result = await streamText({
       messages,
@@ -134,13 +136,27 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       contextOptimization,
     });
 
+    (async () => {
+      for await (const part of result.fullStream) {
+        if (part.type === 'error') {
+          const error: any = part.error;
+          logger.error(`${error}`);
+
+          return;
+        }
+      }
+    })();
+
     stream.switchSource(result.toDataStream());
 
+    // return createrespo
     return new Response(stream.readable, {
       status: 200,
       headers: {
-        contentType: 'text/event-stream',
-        connection: 'keep-alive',
+        'Content-Type': 'text/event-stream; charset=utf-8',
+        Connection: 'keep-alive',
+        'Cache-Control': 'no-cache',
+        'Text-Encoding': 'chunked',
       },
     });
   } catch (error: any) {

diff --git a/app/routes/api.check-env-key.ts b/app/routes/api.check-env-key.ts
@@ -1,4 +1,4 @@
-import type { LoaderFunction } from '@remix-run/node';
+import type { LoaderFunction } from '@remix-run/cloudflare';
 import { providerBaseUrlEnvKeys } from '~/utils/constants';
 
 export const loader: LoaderFunction = async ({ context, request }) => {

diff --git a/app/routes/api.enhancer.ts b/app/routes/api.enhancer.ts
@@ -107,7 +107,10 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
     return new Response(result.textStream, {
       status: 200,
       headers: {
-        'Content-Type': 'text/plain; charset=utf-8',
+        'Content-Type': 'text/event-stream',
+        Connection: 'keep-alive',
+        'Cache-Control': 'no-cache',
+        'Text-Encoding': 'chunked',
       },
     });
   } catch (error: unknown) {