@proteinjs/conversation 1.7.4 → 2.0.0

package/src/OpenAiStreamProcessor.ts

@@ -2,6 +2,7 @@ import { ChatCompletionMessageToolCall, ChatCompletionChunk } from 'openai/resou
 import { LogLevel, Logger } from '@proteinjs/logger';
 import { Stream } from 'openai/streaming';
 import { Readable, Transform, TransformCallback, PassThrough } from 'stream';
+import { UsageData, UsageDataAccumulator } from './UsageData';
 
 export interface AssistantResponseStreamChunk {
   content?: string;
@@ -21,6 +22,7 @@ export class OpenAiStreamProcessor {
   private inputStream: Readable;
   private controlStream: Transform;
   private outputStream: Readable;
+  private outputStreamTerminated = false;
 
   constructor(
     inputStream: Stream<ChatCompletionChunk>,
@@ -28,8 +30,10 @@ export class OpenAiStreamProcessor {
       toolCalls: ChatCompletionMessageToolCall[],
       currentFunctionCalls: number
     ) => Promise<Readable>,
-    logLevel: LogLevel,
-    private abortSignal?: AbortSignal
+    private usageDataAccumulator: UsageDataAccumulator,
+    logLevel?: LogLevel,
+    private abortSignal?: AbortSignal,
+    private onUsageData?: (usageData: UsageData) => Promise<void>
   ) {
     this.logger = new Logger({ name: this.constructor.name, logLevel });
     this.inputStream = Readable.from(inputStream);
@@ -50,6 +54,7 @@ export class OpenAiStreamProcessor {
    * @returns a `Transform` that parses the input stream and delegates to tool calls or writes a text response to the user
    */
   private createControlStream(): Transform {
+    let finishedProcessingToolCallStream = false;
     return new Transform({
       objectMode: true,
       transform: (chunk: ChatCompletionChunk, encoding: string, callback: TransformCallback) => {
@@ -68,21 +73,35 @@ export class OpenAiStreamProcessor {
           } else if (chunk.choices[0]?.delta?.tool_calls) {
             this.handleToolCallDelta(chunk.choices[0].delta.tool_calls);
           } else if (chunk.choices[0]?.finish_reason === 'tool_calls') {
-            this.handleToolCalls();
+            finishedProcessingToolCallStream = true;
           } else if (chunk.choices[0]?.finish_reason === 'stop') {
             this.outputStream.push({ finishReason: 'stop' } as AssistantResponseStreamChunk);
             this.outputStream.push(null);
-            this.destroyStreams();
+            this.outputStreamTerminated = true;
           } else if (chunk.choices[0]?.finish_reason === 'length') {
             this.logger.info({ message: `The maximum number of output tokens was reached` });
             this.outputStream.push({ finishReason: 'length' } as AssistantResponseStreamChunk);
             this.outputStream.push(null);
-            this.destroyStreams();
+            this.outputStreamTerminated = true;
           } else if (chunk.choices[0]?.finish_reason === 'content_filter') {
             this.logger.warn({ message: `Content was omitted due to a flag from OpenAI's content filters` });
             this.outputStream.push({ finishReason: 'content_filter' } as AssistantResponseStreamChunk);
             this.outputStream.push(null);
-            this.destroyStreams();
+            this.outputStreamTerminated = true;
+          } else if (chunk.usage) {
+            this.usageDataAccumulator.addTokenUsage({
+              promptTokens: chunk.usage.prompt_tokens,
+              completionTokens: chunk.usage.completion_tokens,
+              totalTokens: chunk.usage.total_tokens,
+            });
+            if (finishedProcessingToolCallStream) {
+              this.handleToolCalls();
+            } else if (this.outputStreamTerminated) {
+              if (this.onUsageData) {
+                this.onUsageData(this.usageDataAccumulator.usageData);
+              }
+              this.destroyStreams();
+            }
           }
           callback();
         } catch (error: any) {

package/src/UsageData.ts

@@ -0,0 +1,76 @@
+import { TiktokenModel } from 'tiktoken';
+
+export type TokenUsage = {
+  promptTokens: number;
+  completionTokens: number;
+  totalTokens: number;
+};
+
+/**
+ * Usage data accumulated throughout the lifecycle of a single call to
+ * `OpenAi.generateResponse` or `OpenAi.generateStreamingResponse`.
+ */
+export type UsageData = {
+  /** The model used by the assistant */
+  model: TiktokenModel;
+  /** The token usage of the initial request sent to the assistant */
+  initialRequestTokenUsage: TokenUsage;
+  /** The total token usage of all requests sent to the assistant (ie. initial request + all subsequent tool call requests) */
+  totalTokenUsage: TokenUsage;
+  /** The number of requests sent to the assistant */
+  totalRequestsToAssistant: number;
+  /** The number of times each tool was called by the assistant */
+  callsPerTool: { [toolName: string]: number };
+  /** The total number of tool calls made by the assistant */
+  totalToolCalls: number;
+};
+
+type UsageDataAccumulatorParams = {
+  model: TiktokenModel;
+};
+
+export class UsageDataAccumulator {
+  private processedInitialRequest = false;
+  public usageData: UsageData;
+
+  constructor({ model }: UsageDataAccumulatorParams) {
+    this.usageData = {
+      model,
+      initialRequestTokenUsage: {
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0,
+      },
+      totalTokenUsage: {
+        promptTokens: 0,
+        completionTokens: 0,
+        totalTokens: 0,
+      },
+      totalRequestsToAssistant: 0,
+      callsPerTool: {},
+      totalToolCalls: 0,
+    };
+  }
+
+  addTokenUsage(tokenUsage: TokenUsage) {
+    this.usageData.totalRequestsToAssistant++;
+    if (!this.processedInitialRequest) {
+      this.usageData.initialRequestTokenUsage = tokenUsage;
+      this.processedInitialRequest = true;
+    }
+    this.usageData.totalTokenUsage = {
+      promptTokens: this.usageData.totalTokenUsage.promptTokens + tokenUsage.promptTokens,
+      completionTokens: this.usageData.totalTokenUsage.completionTokens + tokenUsage.completionTokens,
+      totalTokens: this.usageData.totalTokenUsage.totalTokens + tokenUsage.totalTokens,
+    };
+  }
+
+  recordToolCall(toolName: string) {
+    if (!this.usageData.callsPerTool[toolName]) {
+      this.usageData.callsPerTool[toolName] = 0;
+    }
+
+    this.usageData.callsPerTool[toolName]++;
+    this.usageData.totalToolCalls++;
+  }
+}

package/src/code_template/Code.ts

@@ -31,7 +31,7 @@ export class Code {
       this.addImports(this.args.imports, this.args.conversation);
     }
 
-    return await this.args.conversation.generateCode(this.args.description, 'gpt-4');
+    return await this.args.conversation.generateCode({ description: this.args.description, model: 'gpt-4' });
   }
 
   private addImports(imports: Import[], conversation: Conversation) {

package/test/openai/openai.generateList.test.ts

@@ -2,7 +2,7 @@ import { OpenAi } from '../../src/OpenAi';
 
 test('generateList should return an array of numbers, counting to 10', async () => {
   const numbers = ['one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine', 'ten'];
-  expect((await OpenAi.generateList([`Create a list of numbers spelled out, from 1 to 10`])).join(' ')).toBe(
-    numbers.join(' ')
-  );
+  expect(
+    (await new OpenAi().generateList({ messages: [`Create a list of numbers spelled out, from 1 to 10`] })).join(' ')
+  ).toBe(numbers.join(' '));
 });