Merge remote-tracking branch 'origin/main' into system-prompt-variations-local

app/routes/api.chat.ts

@@ -1,4 +1,5 @@
 import { type ActionFunctionArgs } from '@remix-run/cloudflare';
+import { createDataStream } from 'ai';
 import { MAX_RESPONSE_SEGMENTS, MAX_TOKENS } from '~/lib/.server/llm/constants';
 import { CONTINUE_PROMPT } from '~/lib/common/prompts/prompts';
 import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
@@ -9,17 +10,15 @@ export async function action(args: ActionFunctionArgs) {
   return chatAction(args);
 }
 
-function parseCookies(cookieHeader: string) {
-  const cookies: any = {};
+function parseCookies(cookieHeader: string): Record<string, string> {
+  const cookies: Record<string, string> = {};
 
-  // Split the cookie string by semicolons and spaces
   const items = cookieHeader.split(';').map((cookie) => cookie.trim());
 
   items.forEach((item) => {
     const [name, ...rest] = item.split('=');
 
     if (name && rest) {
-      // Decode the name and value, and join value parts in case it contains '='
       const decodedName = decodeURIComponent(name.trim());
       const decodedValue = decodeURIComponent(rest.join('=').trim());
       cookies[decodedName] = decodedValue;
@@ -37,8 +36,6 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
   }>();
 
   const cookieHeader = request.headers.get('Cookie');
-
-  // Parse the cookie's value (returns an object or null if no cookie exists)
   const apiKeys = JSON.parse(parseCookies(cookieHeader || '').apiKeys || '{}');
   const providerSettings: Record<string, IProviderSetting> = JSON.parse(
     parseCookies(cookieHeader || '').providers || '{}',
@@ -46,12 +43,42 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
 
   const stream = new SwitchableStream();
 
+  const cumulativeUsage = {
+    completionTokens: 0,
+    promptTokens: 0,
+    totalTokens: 0,
+  };
+
   try {
     const options: StreamingOptions = {
       toolChoice: 'none',
-      onFinish: async ({ text: content, finishReason }) => {
+      onFinish: async ({ text: content, finishReason, usage }) => {
+        console.log('usage', usage);
+
+        if (usage) {
+          cumulativeUsage.completionTokens += usage.completionTokens || 0;
+          cumulativeUsage.promptTokens += usage.promptTokens || 0;
+          cumulativeUsage.totalTokens += usage.totalTokens || 0;
+        }
+
         if (finishReason !== 'length') {
-          return stream.close();
+          return stream
+            .switchSource(
+              createDataStream({
+                async execute(dataStream) {
+                  dataStream.writeMessageAnnotation({
+                    type: 'usage',
+                    value: {
+                      completionTokens: cumulativeUsage.completionTokens,
+                      promptTokens: cumulativeUsage.promptTokens,
+                      totalTokens: cumulativeUsage.totalTokens,
+                    },
+                  });
+                },
+                onError: (error: any) => `Custom error: ${error.message}`,
+              }),
+            )
+            .then(() => stream.close());
         }
 
         if (stream.switches >= MAX_RESPONSE_SEGMENTS) {
@@ -75,7 +102,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
           promptId,
         });
 
-        return stream.switchSource(result.toAIStream());
+        return stream.switchSource(result.toDataStream());
       },
     };
 
@@ -89,7 +116,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       promptId,
     });
 
-    stream.switchSource(result.toAIStream());
+    stream.switchSource(result.toDataStream());
 
     return new Response(stream.readable, {
       status: 200,
@@ -98,7 +125,7 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
       },
     });
   } catch (error: any) {
-    console.log(error);
+    console.error(error);
 
     if (error.message?.includes('API key')) {
       throw new Response('Invalid or missing API key', {

app/routes/api.enhancer.ts

@@ -1,5 +1,6 @@
 import { type ActionFunctionArgs } from '@remix-run/cloudflare';
-import { StreamingTextResponse, parseStreamPart } from 'ai';
+
+//import { StreamingTextResponse, parseStreamPart } from 'ai';
 import { streamText } from '~/lib/.server/llm/stream-text';
 import { stripIndents } from '~/utils/stripIndent';
 import type { IProviderSetting, ProviderInfo } from '~/types/model';
@@ -73,32 +74,32 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
             `[Model: ${model}]\n\n[Provider: ${providerName}]\n\n` +
             stripIndents`
             You are a professional prompt engineer specializing in crafting precise, effective prompts.
-          Your task is to enhance prompts by making them more specific, actionable, and effective.
+            Your task is to enhance prompts by making them more specific, actionable, and effective.
 
-          I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
+            I want you to improve the user prompt that is wrapped in \`<original_prompt>\` tags.
 
-          For valid prompts:
-          - Make instructions explicit and unambiguous
-          - Add relevant context and constraints
-          - Remove redundant information
-          - Maintain the core intent
-          - Ensure the prompt is self-contained
-          - Use professional language
+            For valid prompts:
+            - Make instructions explicit and unambiguous
+            - Add relevant context and constraints
+            - Remove redundant information
+            - Maintain the core intent
+            - Ensure the prompt is self-contained
+            - Use professional language
 
-          For invalid or unclear prompts:
-          - Respond with a clear, professional guidance message
-          - Keep responses concise and actionable
-          - Maintain a helpful, constructive tone
-          - Focus on what the user should provide
-          - Use a standard template for consistency
+            For invalid or unclear prompts:
+            - Respond with clear, professional guidance
+            - Keep responses concise and actionable
+            - Maintain a helpful, constructive tone
+            - Focus on what the user should provide
+            - Use a standard template for consistency
 
-          IMPORTANT: Your response must ONLY contain the enhanced prompt text.
-          Do not include any explanations, metadata, or wrapper tags.
+            IMPORTANT: Your response must ONLY contain the enhanced prompt text.
+            Do not include any explanations, metadata, or wrapper tags.
 
-          <original_prompt>
-            ${message}
-          </original_prompt>
-        `,
+            <original_prompt>
+              ${message}
+            </original_prompt>
+          `,
         },
       ],
       env: context.cloudflare.env,
@@ -113,7 +114,7 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
 
         for (const line of lines) {
           try {
-            const parsed = parseStreamPart(line);
+            const parsed = JSON.parse(line);
 
             if (parsed.type === 'text') {
               controller.enqueue(encoder.encode(parsed.value));
@@ -128,7 +129,12 @@ async function enhancerAction({ context, request }: ActionFunctionArgs) {
 
     const transformedStream = result.toDataStream().pipeThrough(transformStream);
 
-    return new StreamingTextResponse(transformedStream);
+    return new Response(transformedStream, {
+      status: 200,
+      headers: {
+        'Content-Type': 'text/plain; charset=utf-8',
+      },
+    });
   } catch (error: unknown) {
     console.log(error);