Async Iterables and Thenable LLM Calls

src/gpt-response.ts

@@ -0,0 +1,157 @@
+import type { SSEvent } from './sse-session.js';
+
+export type MessageChunk = {
+  type: 'reasoning' | 'content';
+  reasoning_details?: string;
+  content: string;
+}
+
+export type FinalResult = {
+  reasoning: string;
+  content: string;
+}
+
+export type GPTResponse = {
+  id: string;
+  provider: string;
+  model: string;
+  object: string;
+  created: number;
+  choices: {
+    index: number;
+    delta: {
+      role: 'user' | 'assistant' | 'system';
+      content: string;
+      reasoning: string;
+      reasoning_details: {
+        type: string;
+        summary: string;
+      }
+    };
+  }[];
+  finish_reason: 'stop' | 'tool_calls' | 'length' | 'content_filter' | null;
+  native_finish_reason: string | null;
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+    cost: number;
+    is_byok: boolean;
+    prompt_tokens_details: {
+      cached_tokens: number;
+    };
+    cost_details: {
+      upstream_inference_cost: number;
+      upstream_prompt_cost: number;
+      upstream_inference_completions_cost: number;
+    },
+    completion_tokens_details: {
+      reasoning_tokens: number;
+    }
+  };
+}
+
+export class MessageResponse implements PromiseLike<FinalResult> {
+  private chunks: MessageChunk[] = [];
+  private iteratorConsumed = false;
+  private resolveResult!: (value: FinalResult) => void;
+  private resultPromise: Promise<FinalResult>;
+  private iterator: AsyncIterable<SSEvent>;
+
+  constructor(iterator: AsyncIterable<SSEvent>) {
+    this.iterator = iterator;
+    this.resultPromise = new Promise(resolve => {
+      this.resolveResult = resolve;
+    });
+  }
+
+  async *[Symbol.asyncIterator]() {
+    if (this.iteratorConsumed) {
+      throw new Error('GPTResponse can only be iterated once');
+    }
+    
+    this.iteratorConsumed = true;
+
+    for await (const rawChunk of this.iterator) {
+      const chunk = this.parseChunk(rawChunk);
+      this.chunks.push(chunk);
+      yield chunk;
+    }
+
+    this.resolveResult(this.buildResult());
+  }
+
+  then<TResult1 = FinalResult, TResult2 = never>(
+    onfulfilled?: ((value: FinalResult) => TResult1 | PromiseLike<TResult1>) | null,
+    onrejected?: ((reason: unknown) => TResult2 | PromiseLike<TResult2>) | null,
+  ): Promise<TResult1 | TResult2> {
+    // If not yet iterated, consume the iterator to get the result
+    if (!this.iteratorConsumed) {
+      this.iteratorConsumed = true;
+
+      (async () => {
+        for await (const rawChunk of this.iterator) {
+          const chunk = this.parseChunk(rawChunk);
+          this.chunks.push(chunk);
+        }
+
+        this.resolveResult(this.buildResult());
+      })();
+    }
+
+    return this.resultPromise.then(onfulfilled, onrejected);
+  }
+
+  catch(onrejected?: ((reason: unknown) => never) | null): Promise<FinalResult> {
+    return this.resultPromise.catch(onrejected);
+  }
+
+  finally(onfinally?: (() => void) | undefined): Promise<FinalResult> {
+    return this.resultPromise.finally(onfinally);
+  }
+
+  private buildResult(): FinalResult {
+    return {
+      reasoning: this.chunks
+        .filter(c => c.type === 'reasoning')
+        .map(c => c.content)
+        .join(''),
+      content: this.chunks
+        .filter(c => c.type === 'content')
+        .map(c => c.content)
+        .join(''),
+    };
+  }
+
+  private parseChunk(rawChunk: SSEvent) {
+    // console.log('Raw Chunk:', rawChunk);
+    if (rawChunk.data === '[DONE]') {
+      return {
+        type: 'content',
+        content: '',
+      } as const;
+    }
+    
+    const data = JSON.parse(rawChunk.data) as GPTResponse;
+    const choice = data.choices[0];
+
+    if (!choice) {
+      throw new Error('No choice found in chunk');
+    }
+    
+    const delta = choice.delta;
+    
+    if (delta.reasoning) {
+      return {
+        type: 'reasoning',
+        content: delta.reasoning,
+        reasoning_details: delta.reasoning_details.summary,
+      } as const;
+    } else {
+      return {
+        type: 'content',
+        content: delta.content,
+      } as const;
+    }
+  }
+}