@lobehub/chat 1.68.9 → 1.68.11

package/src/libs/agent-runtime/togetherai/index.test.ts

@@ -1,300 +1,12 @@
 // @vitest-environment node
-import OpenAI from 'openai';
-import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { testProvider } from '@/libs/agent-runtime/providerTestUtils';
 
-import { LobeOpenAICompatibleRuntime } from '@/libs/agent-runtime';
-
-import * as debugStreamModule from '../utils/debugStream';
-import models from './fixtures/models.json';
 import { LobeTogetherAI } from './index';
 
-const provider = 'togetherai';
-const defaultBaseURL = 'https://api.together.xyz/v1';
-
-const bizErrorType = 'ProviderBizError';
-const invalidErrorType = 'InvalidProviderAPIKey';
-
-// Mock the console.error to avoid polluting test output
-vi.spyOn(console, 'error').mockImplementation(() => {});
-
-let instance: LobeOpenAICompatibleRuntime;
-
-beforeEach(() => {
-  instance = new LobeTogetherAI({ apiKey: 'test' });
-
-  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
-  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
-    new ReadableStream() as any,
-  );
-});
-
-afterEach(() => {
-  vi.clearAllMocks();
-});
-
-describe('LobeTogetherAI', () => {
-  describe('init', () => {
-    it('should correctly initialize with an API key', async () => {
-      const instance = new LobeTogetherAI({ apiKey: 'test_api_key' });
-      expect(instance).toBeInstanceOf(LobeTogetherAI);
-      expect(instance.baseURL).toEqual(defaultBaseURL);
-    });
-  });
-
-  describe('chat', () => {
-    it('should return a StreamingTextResponse on successful API call', async () => {
-      // Arrange
-      const mockStream = new ReadableStream();
-      const mockResponse = Promise.resolve(mockStream);
-
-      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
-
-      // Act
-      const result = await instance.chat({
-        messages: [{ content: 'Hello', role: 'user' }],
-        model: 'mistralai/mistral-7b-instruct:free',
-        temperature: 0,
-      });
-
-      // Assert
-      expect(result).toBeInstanceOf(Response);
-    });
-
-    it('should call TogetherAI API with corresponding options', async () => {
-      // Arrange
-      const mockStream = new ReadableStream();
-      const mockResponse = Promise.resolve(mockStream);
-
-      (instance['client'].chat.completions.create as Mock).mockResolvedValue(mockResponse);
-
-      // Act
-      const result = await instance.chat({
-        max_tokens: 1024,
-        messages: [{ content: 'Hello', role: 'user' }],
-        model: 'mistralai/mistral-7b-instruct:free',
-        temperature: 0.7,
-        top_p: 1,
-      });
-
-      // Assert
-      expect(instance['client'].chat.completions.create).toHaveBeenCalledWith(
-        {
-          max_tokens: 1024,
-          messages: [{ content: 'Hello', role: 'user' }],
-          model: 'mistralai/mistral-7b-instruct:free',
-          temperature: 0.7,
-          stream: true,
-          top_p: 1,
-        },
-        { headers: { Accept: '*/*' } },
-      );
-      expect(result).toBeInstanceOf(Response);
-    });
-
-    describe('Error', () => {
-      it('should return TogetherAIBizError with an openai error response when OpenAI.APIError is thrown', async () => {
-        // Arrange
-        const apiError = new OpenAI.APIError(
-          400,
-          {
-            status: 400,
-            error: {
-              message: 'Bad Request',
-            },
-          },
-          'Error message',
-          {},
-        );
-
-        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
-
-        // Act
-        try {
-          await instance.chat({
-            messages: [{ content: 'Hello', role: 'user' }],
-            model: 'mistralai/mistral-7b-instruct:free',
-            temperature: 0,
-          });
-        } catch (e) {
-          expect(e).toEqual({
-            endpoint: defaultBaseURL,
-            error: {
-              error: { message: 'Bad Request' },
-              status: 400,
-            },
-            errorType: bizErrorType,
-            provider,
-          });
-        }
-      });
-
-      it('should throw AgentRuntimeError with InvalidTogetherAIAPIKey if no apiKey is provided', async () => {
-        try {
-          new LobeTogetherAI({});
-        } catch (e) {
-          expect(e).toEqual({ errorType: invalidErrorType });
-        }
-      });
-
-      it('should return TogetherAIBizError with the cause when OpenAI.APIError is thrown with cause', async () => {
-        // Arrange
-        const errorInfo = {
-          stack: 'abc',
-          cause: {
-            message: 'api is undefined',
-          },
-        };
-        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
-
-        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
-
-        // Act
-        try {
-          await instance.chat({
-            messages: [{ content: 'Hello', role: 'user' }],
-            model: 'mistralai/mistral-7b-instruct:free',
-            temperature: 0,
-          });
-        } catch (e) {
-          expect(e).toEqual({
-            endpoint: defaultBaseURL,
-            error: {
-              cause: { message: 'api is undefined' },
-              stack: 'abc',
-            },
-            errorType: bizErrorType,
-            provider,
-          });
-        }
-      });
-
-      it('should return TogetherAIBizError with an cause response with desensitize Url', async () => {
-        // Arrange
-        const errorInfo = {
-          stack: 'abc',
-          cause: { message: 'api is undefined' },
-        };
-        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
-
-        instance = new LobeTogetherAI({
-          apiKey: 'test',
-
-          baseURL: 'https://api.abc.com/v1',
-        });
-
-        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
-
-        // Act
-        try {
-          await instance.chat({
-            messages: [{ content: 'Hello', role: 'user' }],
-            model: 'mistralai/mistral-7b-instruct:free',
-            temperature: 0,
-          });
-        } catch (e) {
-          expect(e).toEqual({
-            endpoint: 'https://api.***.com/v1',
-            error: {
-              cause: { message: 'api is undefined' },
-              stack: 'abc',
-            },
-            errorType: bizErrorType,
-            provider,
-          });
-        }
-      });
-
-      it('should throw an InvalidTogetherAIAPIKey error type on 401 status code', async () => {
-        // Mock the API call to simulate a 401 error
-        const error = new Error('Unauthorized') as any;
-        error.status = 401;
-        vi.mocked(instance['client'].chat.completions.create).mockRejectedValue(error);
-
-        try {
-          await instance.chat({
-            messages: [{ content: 'Hello', role: 'user' }],
-            model: 'mistralai/mistral-7b-instruct:free',
-            temperature: 0,
-          });
-        } catch (e) {
-          // Expect the chat method to throw an error with InvalidTogetherAIAPIKey
-          expect(e).toEqual({
-            endpoint: defaultBaseURL,
-            error: new Error('Unauthorized'),
-            errorType: invalidErrorType,
-            provider,
-          });
-        }
-      });
-
-      it('should return AgentRuntimeError for non-OpenAI errors', async () => {
-        // Arrange
-        const genericError = new Error('Generic Error');
-
-        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(genericError);
-
-        // Act
-        try {
-          await instance.chat({
-            messages: [{ content: 'Hello', role: 'user' }],
-            model: 'mistralai/mistral-7b-instruct:free',
-            temperature: 0,
-          });
-        } catch (e) {
-          expect(e).toEqual({
-            endpoint: defaultBaseURL,
-            errorType: 'AgentRuntimeError',
-            provider,
-            error: {
-              name: genericError.name,
-              cause: genericError.cause,
-              message: genericError.message,
-              stack: genericError.stack,
-            },
-          });
-        }
-      });
-    });
-
-    describe('DEBUG', () => {
-      it('should call debugStream and return StreamingTextResponse when DEBUG_TOGETHERAI_CHAT_COMPLETION is 1', async () => {
-        // Arrange
-        const mockProdStream = new ReadableStream() as any; // 模拟的 prod 流
-        const mockDebugStream = new ReadableStream({
-          start(controller) {
-            controller.enqueue('Debug stream content');
-            controller.close();
-          },
-        }) as any;
-        mockDebugStream.toReadableStream = () => mockDebugStream; // 添加 toReadableStream 方法
-
-        // 模拟 chat.completions.create 返回值，包括模拟的 tee 方法
-        (instance['client'].chat.completions.create as Mock).mockResolvedValue({
-          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
-        });
-
-        // 保存原始环境变量值
-        const originalDebugValue = process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION;
-
-        // 模拟环境变量
-        process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION = '1';
-        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
-
-        // 执行测试
-        // 运行你的测试函数，确保它会在条件满足时调用 debugStream
-        // 假设的测试函数调用，你可能需要根据实际情况调整
-        await instance.chat({
-          messages: [{ content: 'Hello', role: 'user' }],
-          model: 'mistralai/mistral-7b-instruct:free',
-          temperature: 0,
-        });
-
-        // 验证 debugStream 被调用
-        expect(debugStreamModule.debugStream).toHaveBeenCalled();
-
-        // 恢复原始环境变量值
-        process.env.DEBUG_TOGETHERAI_CHAT_COMPLETION = originalDebugValue;
-      });
-    });
-  });
+testProvider({
+  provider: 'togetherai',
+  defaultBaseURL: 'https://api.together.xyz/v1',
+  chatModel: 'mistralai/mistral-7b-instruct:free',
+  Runtime: LobeTogetherAI,
+  chatDebugEnv: 'DEBUG_TOGETHERAI_CHAT_COMPLETION',
 });

package/src/libs/agent-runtime/utils/openaiCompatibleFactory/index.test.ts

@@ -106,6 +106,9 @@ describe('LobeOpenAICompatibleFactory', () => {
           model: 'mistralai/mistral-7b-instruct:free',
           temperature: 0.7,
           stream: true,
+          stream_options: {
+            include_usage: true,
+          },
           top_p: 1,
         },
         { headers: { Accept: '*/*' } },

package/src/libs/agent-runtime/utils/openaiCompatibleFactory/index.ts

@@ -101,8 +101,9 @@ interface OpenAICompatibleFactoryOptions<T extends Record<string, any> = any> {
 export function transformResponseToStream(data: OpenAI.ChatCompletion) {
   return new ReadableStream({
     start(controller) {
+      const choices = data.choices || [];
       const chunk: OpenAI.ChatCompletionChunk = {
-        choices: data.choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
+        choices: choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
           delta: {
             content: choice.message.content,
             role: choice.message.role,
@@ -128,7 +129,7 @@ export function transformResponseToStream(data: OpenAI.ChatCompletion) {
       controller.enqueue(chunk);
 
       controller.enqueue({
-        choices: data.choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
+        choices: choices.map((choice: OpenAI.ChatCompletion.Choice) => ({
           delta: {
             content: null,
             role: choice.message.role,
@@ -219,7 +220,9 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
             ...postPayload,
             messages,
             ...(chatCompletion?.noUserId ? {} : { user: options?.user }),
+            stream_options: postPayload.stream ? { include_usage: true } : undefined,
           };
+
           if (debug?.chatCompletion?.()) {
             console.log('[requestPayload]:', JSON.stringify(finalPayload, null, 2));
           }

package/src/libs/agent-runtime/utils/streams/anthropic.test.ts

@@ -225,7 +225,7 @@ describe('AnthropicStream', () => {
 
         'id: msg_017aTuY86wNxth5TE544yqJq',
         'event: usage',
-        'data: {"inputTokens":457,"outputTokens":84,"totalTokens":541}\n',
+        'data: {"inputCacheMissTokens":457,"totalInputTokens":457,"totalOutputTokens":84,"totalTokens":541}\n',
       ].map((item) => `${item}\n`),
     );
 
@@ -381,8 +381,7 @@ describe('AnthropicStream', () => {
 
         'id: msg_0175ryA67RbGrnRrGBXFQEYK',
         'event: usage',
-        'data: {"inputTokens":485,"outputTokens":154,"totalTokens":639}\n',
-
+        'data: {"inputCacheMissTokens":485,"totalInputTokens":485,"totalOutputTokens":154,"totalTokens":639}\n',
         'id: msg_0175ryA67RbGrnRrGBXFQEYK',
         'event: stop',
         'data: "message_stop"\n',
@@ -392,6 +391,91 @@ describe('AnthropicStream', () => {
     expect(onToolCallMock).toHaveBeenCalledTimes(6);
   });
 
+  it('should handle prompts context caching', async () => {
+    const streams = [
+      {
+        type: 'message_start',
+        message: {
+          id: 'msg_01Vxc4yQTEjkDSba3N3BMbH8',
+          type: 'message',
+          role: 'assistant',
+          model: 'claude-3-7-sonnet-20250219',
+          content: [],
+          stop_reason: null,
+          stop_sequence: null,
+          usage: {
+            input_tokens: 6,
+            cache_creation_input_tokens: 457,
+            cache_read_input_tokens: 17918,
+            output_tokens: 2,
+          },
+        },
+      },
+      { type: 'content_block_start', index: 0, content_block: { type: 'text', text: '' } },
+      { type: 'content_block_delta', index: 0, delta: { type: 'text_delta', text: '\n\n根' } },
+      {
+        type: 'content_block_delta',
+        index: 0,
+        delta: { type: 'text_delta', text: '/\n[^20]: https://s' },
+      },
+      { type: 'content_block_stop', index: 0 },
+      {
+        type: 'message_delta',
+        delta: { stop_reason: 'end_turn', stop_sequence: null },
+        usage: { output_tokens: 3222 },
+      },
+      { type: 'message_stop' },
+    ];
+
+    const mockReadableStream = new ReadableStream({
+      start(controller) {
+        streams.forEach((chunk) => {
+          controller.enqueue(chunk);
+        });
+        controller.close();
+      },
+    });
+
+    const protocolStream = AnthropicStream(mockReadableStream);
+
+    const decoder = new TextDecoder();
+    const chunks = [];
+
+    // @ts-ignore
+    for await (const chunk of protocolStream) {
+      chunks.push(decoder.decode(chunk, { stream: true }));
+    }
+
+    expect(chunks).toEqual(
+      [
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: data',
+        'data: {"id":"msg_01Vxc4yQTEjkDSba3N3BMbH8","type":"message","role":"assistant","model":"claude-3-7-sonnet-20250219","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":6,"cache_creation_input_tokens":457,"cache_read_input_tokens":17918,"output_tokens":2}}\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: data',
+        'data: ""\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: text',
+        'data: "\\n\\n根"\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: text',
+        'data: "/\\n[^20]: https://s"\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: data',
+        'data: {"type":"content_block_stop","index":0}\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: stop',
+        'data: "end_turn"\n',
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: usage',
+        'data: {"inputCacheMissTokens":6,"inputCachedTokens":17918,"inputWriteCacheTokens":457,"totalInputTokens":18381,"totalOutputTokens":3224,"totalTokens":21605}\n',
+
+        'id: msg_01Vxc4yQTEjkDSba3N3BMbH8',
+        'event: stop',
+        'data: "message_stop"\n',
+      ].map((item) => `${item}\n`),
+    );
+  });
   describe('thinking', () => {
     it('should handle normal thinking ', async () => {
       const streams = [
@@ -515,7 +599,7 @@ describe('AnthropicStream', () => {
           'data: "end_turn"\n',
           'id: msg_01MNsLe7n1uVLtu6W8rCFujD',
           'event: usage',
-          'data: {"inputTokens":46,"outputTokens":365,"totalTokens":411}\n',
+          'data: {"inputCacheMissTokens":46,"totalInputTokens":46,"totalOutputTokens":365,"totalTokens":411}\n',
           'id: msg_01MNsLe7n1uVLtu6W8rCFujD',
           'event: stop',
           'data: "message_stop"\n',
@@ -675,7 +759,7 @@ describe('AnthropicStream', () => {
           'data: "end_turn"\n',
           'id: msg_019q32esPvu3TftzZnL6JPys',
           'event: usage',
-          'data: {"inputTokens":92,"outputTokens":263,"totalTokens":355}\n',
+          'data: {"inputCacheMissTokens":92,"totalInputTokens":92,"totalOutputTokens":263,"totalTokens":355}\n',
           'id: msg_019q32esPvu3TftzZnL6JPys',
           'event: stop',
           'data: "message_stop"\n',

package/src/libs/agent-runtime/utils/streams/anthropic.ts

@@ -22,9 +22,24 @@ export const transformAnthropicStream = (
   switch (chunk.type) {
     case 'message_start': {
       context.id = chunk.message.id;
+      let totalInputTokens = chunk.message.usage?.input_tokens;
+
+      if (
+        chunk.message.usage?.cache_creation_input_tokens ||
+        chunk.message.usage?.cache_read_input_tokens
+      ) {
+        totalInputTokens =
+          chunk.message.usage?.input_tokens +
+          (chunk.message.usage.cache_creation_input_tokens || 0) +
+          (chunk.message.usage.cache_read_input_tokens || 0);
+      }
+
       context.usage = {
-        inputTokens: chunk.message.usage?.input_tokens,
-        outputTokens: chunk.message.usage?.output_tokens,
+        inputCacheMissTokens: chunk.message.usage?.input_tokens,
+        inputCachedTokens: chunk.message.usage?.cache_read_input_tokens || undefined,
+        inputWriteCacheTokens: chunk.message.usage?.cache_creation_input_tokens || undefined,
+        totalInputTokens,
+        totalOutputTokens: chunk.message.usage?.output_tokens,
       };
 
       return { data: chunk.message, id: chunk.message.id, type: 'data' };
@@ -140,18 +155,20 @@ export const transformAnthropicStream = (
     }
 
     case 'message_delta': {
-      const outputTokens = chunk.usage?.output_tokens + (context.usage?.outputTokens || 0);
-      const inputTokens = context.usage?.inputTokens || 0;
-      const totalTokens = inputTokens + outputTokens;
+      const totalOutputTokens =
+        chunk.usage?.output_tokens + (context.usage?.totalOutputTokens || 0);
+      const totalInputTokens = context.usage?.totalInputTokens || 0;
+      const totalTokens = totalInputTokens + totalOutputTokens;
 
       if (totalTokens > 0) {
         return [
           { data: chunk.delta.stop_reason, id: context.id, type: 'stop' },
           {
             data: {
-              inputTokens: inputTokens,
-              outputTokens: outputTokens,
-              totalTokens: inputTokens + outputTokens,
+              ...context.usage,
+              totalInputTokens,
+              totalOutputTokens,
+              totalTokens,
             } as ModelTokensUsage,
             id: context.id,
             type: 'usage',