stamp: creating result types and common usage interface

- Improving typescript with conditional output types based on the
selected provider
- Defining common properties for LLM providers like `usage` metrics and
simplified `value`

Author: kallebysantos

ext/ai/js/ai.ts

@@ -34,15 +34,23 @@ export type EmbeddingInputOptions = {
   normalize?: boolean;
 };
 
-export type SessionInputOptions<T extends SessionType> = T extends
-  LLMProviderName ? LLMInputOptions
-  : EmbeddingInputOptions;
+export type SessionInputOptions<T extends SessionType> = T extends "gte-small"
+  ? EmbeddingInputOptions
+  : T extends LLMProviderName ? LLMInputOptions
+  : never;
+
+export type SessionOutput<T extends SessionType, O> = T extends "gte-small"
+  ? number[]
+  : T extends LLMProviderName
+    ? O extends { stream: true }
+      ? AsyncGenerator<LLMProviderInstance<T>["output"]>
+    : LLMProviderInstance<T>["output"]
+  : never;
 
 export class Session<T extends SessionType> {
   #model?: string;
   #init?: Promise<void>;
 
-  // TODO:(kallebysantos) get 'provider' type here and use type checking to suggest Inputs when run
   constructor(
     public readonly type: T,
     public readonly options?: SessionOptions<T>,
@@ -77,7 +85,10 @@ export class Session<T extends SessionType> {
   }
 
   //  /** @param {string | object} prompt Either a String (ollama) or an OpenAI chat completion body object (openaicompatible): https://platform.openai.com/docs/api-reference/chat/create */
-  async run(input: SessionInput<T>, options: SessionInputOptions<T>) {
+  async run<O extends SessionInputOptions<T>>(
+    input: SessionInput<T>,
+    options: O,
+  ): Promise<SessionOutput<T, O>> {
     if (this.isLLMType()) {
       const opts = options as LLMInputOptions;
       const stream = opts.stream ?? false;
@@ -93,7 +104,7 @@ export class Session<T extends SessionType> {
         stream,
         signal: opts.signal,
         timeout: opts.timeout,
-      });
+      }) as SessionOutput<T, typeof options>;
     }
 
     if (this.#init) {

ext/ai/js/llm/llm_session.ts

@@ -34,14 +34,27 @@ export interface ILLMProviderOptions {
 export type ILLMProviderInput<T = string | object> = T extends string ? string
   : T;
 
+export interface ILLMProviderOutput<T = object> {
+  value?: string;
+  usage: {
+    inputTokens: number;
+    outputTokens: number;
+    totalTokens: number;
+  };
+  inner: T;
+}
+
 export interface ILLMProvider {
   // TODO:(kallebysantos) remove 'any'
   // TODO: (kallebysantos) standardised output format
   getStream(
     input: ILLMProviderInput,
     signal: AbortSignal,
-  ): Promise<AsyncIterable<any>>;
-  getText(input: ILLMProviderInput, signal: AbortSignal): Promise<any>;
+  ): Promise<AsyncIterable<ILLMProviderOutput>>;
+  getText(
+    input: ILLMProviderInput,
+    signal: AbortSignal,
+  ): Promise<ILLMProviderOutput>;
 }
 
 export const providers = {
@@ -92,7 +105,7 @@ export class LLMSession {
   run(
     input: ILLMProviderInput,
     opts: LLMSessionRunInputOptions,
-  ): Promise<AsyncIterable<any>> | Promise<any> {
+  ): Promise<AsyncIterable<ILLMProviderOutput>> | Promise<ILLMProviderOutput> {
     const isStream = opts.stream ?? false;
 
     const timeoutSeconds = typeof opts.timeout === "number" ? opts.timeout : 60;

ext/ai/js/llm/providers/ollama.ts

@@ -3,11 +3,13 @@ import {
   ILLMProviderInput,
   ILLMProviderMeta,
   ILLMProviderOptions,
-} from '../llm_session.ts';
-import { parseJSON } from '../utils/json_parser.ts';
+  ILLMProviderOutput,
+} from "../llm_session.ts";
+import { parseJSON } from "../utils/json_parser.ts";
 
 export type OllamaProviderOptions = ILLMProviderOptions;
 export type OllamaProviderInput = ILLMProviderInput<string>;
+export type OllamaProviderOutput = ILLMProviderOutput<OllamaMessage>;
 
 export type OllamaMessage = {
   model: string;
@@ -25,7 +27,7 @@ export type OllamaMessage = {
 
 export class OllamaLLMSession implements ILLMProvider, ILLMProviderMeta {
   input!: OllamaProviderInput;
-  output!: unknown;
+  output!: OllamaProviderOutput;
   options: OllamaProviderOptions;
 
   constructor(opts: OllamaProviderOptions) {
@@ -36,31 +38,34 @@ export class OllamaLLMSession implements ILLMProvider, ILLMProviderMeta {
   async getStream(
     prompt: OllamaProviderInput,
     signal: AbortSignal,
-  ): Promise<AsyncIterable<OllamaMessage>> {
+  ): Promise<AsyncIterable<OllamaProviderOutput>> {
     const generator = await this.generate(
       prompt,
       signal,
       true,
     ) as AsyncGenerator<OllamaMessage>;
 
+    const parser = this.parse;
+
     const stream = async function* () {
       for await (const message of generator) {
-        if ('error' in message) {
+        if ("error" in message) {
           if (message.error instanceof Error) {
             throw message.error;
           } else {
             throw new Error(message.error as string);
           }
         }
 
-        yield message;
+        yield parser(message);
+
         if (message.done) {
           return;
         }
       }
 
       throw new Error(
-        'Did not receive done or success response in stream.',
+        "Did not receive done or success response in stream.",
       );
     };
 
@@ -70,14 +75,28 @@ export class OllamaLLMSession implements ILLMProvider, ILLMProviderMeta {
   async getText(
     prompt: OllamaProviderInput,
     signal: AbortSignal,
-  ): Promise<OllamaMessage> {
+  ): Promise<OllamaProviderOutput> {
     const response = await this.generate(prompt, signal) as OllamaMessage;
 
     if (!response?.done) {
-      throw new Error('Expected a completed response.');
+      throw new Error("Expected a completed response.");
     }
 
-    return response;
+    return this.parse(response);
+  }
+
+  private parse(message: OllamaMessage): OllamaProviderOutput {
+    const { response, prompt_eval_count, eval_count } = message;
+
+    return {
+      value: response,
+      inner: message,
+      usage: {
+        inputTokens: prompt_eval_count,
+        outputTokens: eval_count,
+        totalTokens: prompt_eval_count + eval_count,
+      },
+    };
   }
 
   private async generate(
@@ -86,11 +105,11 @@ export class OllamaLLMSession implements ILLMProvider, ILLMProviderMeta {
     stream: boolean = false,
   ) {
     const res = await fetch(
-      new URL('/api/generate', this.options.baseURL),
+      new URL("/api/generate", this.options.baseURL),
       {
-        method: 'POST',
+        method: "POST",
         headers: {
-          'Content-Type': 'application/json',
+          "Content-Type": "application/json",
         },
         body: JSON.stringify({
           model: this.options.model,
@@ -108,7 +127,7 @@ export class OllamaLLMSession implements ILLMProvider, ILLMProviderMeta {
     }
 
     if (!res.body) {
-      throw new Error('Missing body');
+      throw new Error("Missing body");
     }
 
     if (stream) {

ext/ai/js/llm/providers/openai.ts

@@ -3,6 +3,7 @@ import {
   ILLMProviderInput,
   ILLMProviderMeta,
   ILLMProviderOptions,
+  ILLMProviderOutput,
 } from "../llm_session.ts";
 import { parseJSONOverEventStream } from "../utils/json_parser.ts";
 
@@ -97,11 +98,11 @@ export type OpenAIResponse = {
 export type OpenAICompatibleInput = Omit<OpenAIRequest, "stream" | "model">;
 
 export type OpenAIProviderInput = ILLMProviderInput<OpenAICompatibleInput>;
+export type OpenAIProviderOutput = ILLMProviderOutput<OpenAIResponse>;
 
 export class OpenAILLMSession implements ILLMProvider, ILLMProviderMeta {
   input!: OpenAIProviderInput;
-  // TODO:(kallebysantos) add output types
-  output: unknown;
+  output!: OpenAIProviderOutput;
   options: OpenAIProviderOptions;
 
   constructor(opts: OpenAIProviderOptions) {
@@ -111,13 +112,14 @@ export class OpenAILLMSession implements ILLMProvider, ILLMProviderMeta {
   async getStream(
     prompt: OpenAIProviderInput,
     signal: AbortSignal,
-  ): Promise<AsyncIterable<OpenAIResponse>> {
+  ): Promise<AsyncIterable<OpenAIProviderOutput>> {
     const generator = await this.generate(
       prompt,
       signal,
       true,
     ) as AsyncGenerator<any>; // TODO:(kallebysantos) remove any
 
+    const parser = this.parse;
     const stream = async function* () {
       for await (const message of generator) {
         // TODO:(kallebysantos) Simplify duplicated code for stream error checking
@@ -129,7 +131,7 @@ export class OpenAILLMSession implements ILLMProvider, ILLMProviderMeta {
           }
         }
 
-        yield message;
+        yield parser(message);
         const finishReason = message.choices[0].finish_reason;
 
         if (finishReason) {
@@ -152,7 +154,7 @@ export class OpenAILLMSession implements ILLMProvider, ILLMProviderMeta {
   async getText(
     prompt: OpenAIProviderInput,
     signal: AbortSignal,
-  ): Promise<OpenAIResponse> {
+  ): Promise<OpenAIProviderOutput> {
     const response = await this.generate(
       prompt,
       signal,
@@ -164,9 +166,23 @@ export class OpenAILLMSession implements ILLMProvider, ILLMProviderMeta {
       throw new Error("Expected a completed response.");
     }
 
-    return response;
+    return this.parse(response);
   }
 
+  private parse(message: OpenAIResponse): OpenAIProviderOutput {
+    const { usage } = message;
+
+    return {
+      value: message.choices.at(0)?.message.content ?? undefined,
+      inner: message,
+      usage: {
+        // Usage maybe 'null' while streaming, but the final message will include it
+        inputTokens: usage?.prompt_tokens ?? 0,
+        outputTokens: usage?.completion_tokens ?? 0,
+        totalTokens: usage?.total_tokens ?? 0,
+      },
+    };
+  }
   private async generate(
     input: OpenAICompatibleInput,
     signal: AbortSignal,