Add ApiStream generator to interface with providers

src/api/index.ts

@@ -1,5 +1,4 @@
 import { Anthropic } from "@anthropic-ai/sdk"
-import { Stream } from "@anthropic-ai/sdk/streaming"
 import { ApiConfiguration, ModelInfo } from "../shared/api"
 import { AnthropicHandler } from "./providers/anthropic"
 import { AwsBedrockHandler } from "./providers/bedrock"
@@ -9,15 +8,14 @@ import { OpenAiHandler } from "./providers/openai"
 import { OllamaHandler } from "./providers/ollama"
 import { GeminiHandler } from "./providers/gemini"
 import { OpenAiNativeHandler } from "./providers/openai-native"
-
-export type AnthropicStream = Stream<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
+import { ApiStream } from "./transform/stream"
 
 export interface ApiHandler {
 	createMessage(
 		systemPrompt: string,
 		messages: Anthropic.Messages.MessageParam[],
 		tools: Anthropic.Messages.Tool[]
-	): Promise<AnthropicStream>
+	): ApiStream
 
 	getModel(): { id: string; info: ModelInfo }
 }

src/api/providers/anthropic.ts

@@ -1,5 +1,5 @@
 import { Anthropic } from "@anthropic-ai/sdk"
-import { AnthropicStream, ApiHandler } from "../index"
+import { Stream as AnthropicStream } from "@anthropic-ai/sdk/streaming"
 import {
 	anthropicDefaultModelId,
 	AnthropicModelId,
@@ -7,6 +7,8 @@ import {
 	ApiHandlerOptions,
 	ModelInfo,
 } from "../../shared/api"
+import { ApiHandler } from "../index"
+import { ApiStream } from "../transform/stream"
 
 export class AnthropicHandler implements ApiHandler {
 	private options: ApiHandlerOptions
@@ -20,11 +22,8 @@ export class AnthropicHandler implements ApiHandler {
 		})
 	}
 
-	async createMessage(
-		systemPrompt: string,
-		messages: Anthropic.Messages.MessageParam[],
-		tools: Anthropic.Messages.Tool[]
-	): Promise<AnthropicStream> {
+	async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
+		let stream: AnthropicStream<Anthropic.Beta.PromptCaching.Messages.RawPromptCachingBetaMessageStreamEvent>
 		const modelId = this.getModel().id
 		switch (modelId) {
 			case "claude-3-5-sonnet-20240620":
@@ -39,7 +38,7 @@ export class AnthropicHandler implements ApiHandler {
 				)
 				const lastUserMsgIndex = userMsgIndices[userMsgIndices.length - 1] ?? -1
 				const secondLastMsgUserIndex = userMsgIndices[userMsgIndices.length - 2] ?? -1
-				const stream = this.client.beta.promptCaching.messages.create(
+				stream = await this.client.beta.promptCaching.messages.create(
 					{
 						model: modelId,
 						max_tokens: this.getModel().info.maxTokens,
@@ -92,14 +91,10 @@ export class AnthropicHandler implements ApiHandler {
 						}
 					})()
 				)
-
-				return stream
-
-				// throw new Error("Not implemented")
-				// return { message }
+				break
 			}
 			default: {
-				const stream = await this.client.messages.create({
+				stream = (await this.client.messages.create({
 					model: modelId,
 					max_tokens: this.getModel().info.maxTokens,
 					temperature: 0,
@@ -108,8 +103,67 @@ export class AnthropicHandler implements ApiHandler {
 					// tools,
 					// tool_choice: { type: "auto" },
 					stream: true,
-				})
-				return stream as AnthropicStream
+				})) as any
+				break
+			}
+		}
+
+		for await (const chunk of stream) {
+			switch (chunk.type) {
+				case "message_start":
+					// tells us cache reads/writes/input/output
+					const usage = chunk.message.usage
+					yield {
+						type: "usage",
+						inputTokens: usage.input_tokens || 0,
+						outputTokens: usage.output_tokens || 0,
+						cacheWriteTokens: usage.cache_creation_input_tokens || 0,
+						cacheReadTokens: usage.cache_read_input_tokens || 0,
+					}
+					break
+				case "message_delta":
+					// tells us stop_reason, stop_sequence, and output tokens along the way and at the end of the message
+
+					yield {
+						type: "usage",
+						inputTokens: 0,
+						outputTokens: chunk.usage.output_tokens || 0,
+						cacheWriteTokens: 0,
+						cacheReadTokens: 0,
+					}
+					break
+				case "message_stop":
+					// no usage data, just an indicator that the message is done
+					break
+				case "content_block_start":
+					switch (chunk.content_block.type) {
+						case "text":
+							// we may receive multiple text blocks, in which case just insert a line break between them
+							if (chunk.index > 0) {
+								yield {
+									type: "text",
+									text: "\n",
+								}
+							}
+							yield {
+								type: "text",
+								text: chunk.content_block.text,
+							}
+							break
+					}
+					break
+				case "content_block_delta":
+					switch (chunk.delta.type) {
+						case "text_delta":
+							yield {
+								type: "text",
+								text: chunk.delta.text,
+							}
+							break
+					}
+					break
+				case "content_block_stop":
+					break
 			}
 		}
 	}

src/api/transform/stream.ts

@@ -0,0 +1,16 @@
+export type ApiStream = AsyncGenerator<ApiStreamChunk>
+export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk
+
+export interface ApiStreamTextChunk {
+	type: "text"
+	text: string
+}
+
+export interface ApiStreamUsageChunk {
+	type: "usage"
+	inputTokens: number
+	outputTokens: number
+	cacheWriteTokens: number
+	cacheReadTokens: number
+	totalCost?: number // openrouter
+}