feat: Updated packages/core/src/llm/mistral.ts

packages/core/src/llm/LLM.ts

@@ -10,7 +10,6 @@ import {
 
 import { ChatCompletionMessageParam } from "openai/resources";
 import { LLMOptions } from "portkey-ai";
-import { MessageContent } from "../ChatEngine";
 import { globalsHelper, Tokenizers } from "../GlobalsHelper";
 import {
   ANTHROPIC_AI_PROMPT,
@@ -48,8 +47,9 @@ export interface ChatResponse {
   delta?: string;
 }
 
-// NOTE in case we need CompletionResponse to diverge from ChatResponse in the future
-export type CompletionResponse = ChatResponse;
+export interface CompletionResponse {
+  text: string;
+}
 
 export interface LLMMetadata {
   model: string;
@@ -60,40 +60,57 @@ export interface LLMMetadata {
   tokenizer: Tokenizers | undefined;
 }
 
+export interface LLMChatParamsBase {
+  messages: ChatMessage[];
+  parentEvent?: Event;
+  extraParams?: Record<string, any>;
+}
+
+export interface LLMChatParamsStreaming extends LLMChatParamsBase {
+  stream: true;
+}
+
+export interface LLMChatParamsNonStreaming extends LLMChatParamsBase {
+  stream?: false | null;
+}
+
+export interface LLMCompletionParamsBase {
+  prompt: string;
+  parentEvent?: Event;
+}
+
+export interface LLMCompletionParamsStreaming extends LLMCompletionParamsBase {
+  stream: true;
+}
+
+export interface LLMCompletionParamsNonStreaming
+  extends LLMCompletionParamsBase {
+  stream?: false | null;
+}
+
 /**
  * Unified language model interface
  */
 export interface LLM {
   metadata: LLMMetadata;
-  // Whether a LLM has streaming support
-  hasStreaming: boolean;
   /**
    * Get a chat response from the LLM
-   * @param messages
    *
-   * The return type of chat() and complete() are set by the "streaming" parameter being set to True.
+   * @param params
    */
-  chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    messages: ChatMessage[],
-    parentEvent?: Event,
-    streaming?: T,
-  ): Promise<R>;
+  chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
 
   /**
    * Get a prompt completion from the LLM
-   * @param prompt the prompt to complete
+   * @param params
    */
-  complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    prompt: MessageContent,
-    parentEvent?: Event,
-    streaming?: T,
-  ): Promise<R>;
+  complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
 
   /**
    * Calculates the number of tokens needed for the given chat messages
@@ -101,6 +118,39 @@ export interface LLM {
   tokens(messages: ChatMessage[]): number;
 }
 
+abstract class BaseLLM implements LLM {
+  abstract metadata: LLMMetadata;
+
+  complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse | AsyncIterable<string>> {
+    const { prompt, parentEvent, stream } = params;
+    if (stream) {
+      return this.chat({
+        messages: [{ content: prompt, role: "user" }],
+        parentEvent,
+        stream: true,
+      });
+    }
+    const chatResponse = await this.chat({
+      messages: [{ content: prompt, role: "user" }],
+      parentEvent,
+    });
+    return { text: chatResponse.message.content as string };
+  }
+
+  abstract chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  abstract chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+
+  abstract tokens(messages: ChatMessage[]): number;
+}
+
 export const GPT4_MODELS = {
   "gpt-4": { contextWindow: 8192 },
   "gpt-4-32k": { contextWindow: 32768 },
@@ -125,9 +175,7 @@ export const ALL_AVAILABLE_OPENAI_MODELS = {
 /**
  * OpenAI LLM implementation
  */
-export class OpenAI implements LLM {
-  hasStreaming: boolean = true;
-
+export class OpenAI extends BaseLLM {
   // Per completion OpenAI params
   model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS;
   temperature: number;
@@ -135,7 +183,7 @@ export class OpenAI implements LLM {
   maxTokens?: number;
   additionalChatOptions?: Omit<
     Partial<OpenAILLM.Chat.ChatCompletionCreateParams>,
-    "max_tokens" | "messages" | "model" | "temperature" | "top_p" | "streaming"
+    "max_tokens" | "messages" | "model" | "temperature" | "top_p" | "stream"
   >;
 
   // OpenAI session params
@@ -155,6 +203,7 @@ export class OpenAI implements LLM {
       azure?: AzureOpenAIConfig;
     },
   ) {
+    super();
     this.model = init?.model ?? "gpt-3.5-turbo";
     this.temperature = init?.temperature ?? 0.1;
     this.topP = init?.topP ?? 1;
@@ -247,10 +296,12 @@ export class OpenAI implements LLM {
     }
   }
 
-  async chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(messages: ChatMessage[], parentEvent?: Event, streaming?: T): Promise<R> {
+  chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<string>> {
+    const { messages, parentEvent, stream } = params;
     const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
@@ -266,11 +317,8 @@ export class OpenAI implements LLM {
       ...this.additionalChatOptions,
     };
     // Streaming
-    if (streaming) {
-      if (!this.hasStreaming) {
-        throw Error("No streaming support for this LLM.");
-      }
-      return this.streamChat(messages, parentEvent) as R;
+    if (stream) {
+      return this.streamChat(params);
     }
     // Non-streaming
     const response = await this.session.openai.chat.completions.create({
@@ -281,27 +329,13 @@ export class OpenAI implements LLM {
     const content = response.choices[0].message?.content ?? "";
     return {
       message: { content, role: response.choices[0].message.role },
-    } as R;
+    };
   }
 
-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
-    return this.chat(
-      [{ content: prompt, role: "user" }],
-      parentEvent,
-      streaming,
-    );
-  }
-
-  //We can wrap a stream in a generator to add some additional logging behavior
-  //For future edits: syntax for generator type is <typeof Yield, typeof Return, typeof Accept>
-  //"typeof Accept" refers to what types you'll accept when you manually call generator.next(<AcceptType>)
-  protected async *streamChat(
-    messages: ChatMessage[],
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
+  protected async *streamChat({
+    messages,
+    parentEvent,
+  }: LLMChatParamsStreaming): AsyncIterable<string> {
     const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
       model: this.model,
       temperature: this.temperature,
@@ -360,14 +394,6 @@ export class OpenAI implements LLM {
     }
     return;
   }
-
-  //streamComplete doesn't need to be async because it's child function is already async
-  protected streamComplete(
-    query: string,
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
-    return this.streamChat([{ content: query, role: "user" }], parentEvent);
-  }
 }
 
 export const ALL_AVAILABLE_LLAMADEUCE_MODELS = {
@@ -425,16 +451,16 @@ export enum DeuceChatStrategy {
 /**
  * Llama2 LLM implementation
  */
-export class LlamaDeuce implements LLM {
+export class LlamaDeuce extends BaseLLM {
   model: keyof typeof ALL_AVAILABLE_LLAMADEUCE_MODELS;
   chatStrategy: DeuceChatStrategy;
   temperature: number;
   topP: number;
   maxTokens?: number;
   replicateSession: ReplicateSession;
-  hasStreaming: boolean;
 
   constructor(init?: Partial<LlamaDeuce>) {
+    super();
     this.model = init?.model ?? "Llama-2-70b-chat-4bit";
     this.chatStrategy =
       init?.chatStrategy ??
@@ -447,7 +473,6 @@ export class LlamaDeuce implements LLM {
       init?.maxTokens ??
       ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model].contextWindow; // For Replicate, the default is 500 tokens which is too low.
     this.replicateSession = init?.replicateSession ?? new ReplicateSession();
-    this.hasStreaming = init?.hasStreaming ?? false;
   }
 
   tokens(messages: ChatMessage[]): number {
@@ -592,10 +617,12 @@ If a question does not make any sense, or is not factually coherent, explain why
     };
   }
 
-  async chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(messages: ChatMessage[], _parentEvent?: Event, streaming?: T): Promise<R> {
+  chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<string>> {
+    const { messages, parentEvent, stream } = params;
     const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
       .replicateApi as `${string}/${string}:${string}`;
 
@@ -629,14 +656,7 @@ If a question does not make any sense, or is not factually coherent, explain why
         //^ We need to do this because Replicate returns a list of strings (for streaming functionality which is not exposed by the run function)
         role: "assistant",
       },
-    } as R;
-  }
-
-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
-    return this.chat([{ content: prompt, role: "user" }], parentEvent);
+    };
   }
 }
 
@@ -650,9 +670,7 @@ export const ALL_AVAILABLE_ANTHROPIC_MODELS = {
  * Anthropic LLM implementation
  */
 
-export class Anthropic implements LLM {
-  hasStreaming: boolean = true;
-
+export class Anthropic extends BaseLLM {
   // Per completion Anthropic params
   model: keyof typeof ALL_AVAILABLE_ANTHROPIC_MODELS;
   temperature: number;
@@ -668,6 +686,7 @@ export class Anthropic implements LLM {
   callbackManager?: CallbackManager;
 
   constructor(init?: Partial<Anthropic>) {
+    super();
     this.model = init?.model ?? "claude-2";
     this.temperature = init?.temperature ?? 0.1;
     this.topP = init?.topP ?? 0.999; // Per Ben Mann
@@ -719,20 +738,15 @@ export class Anthropic implements LLM {
     );
   }
 
-  async chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    messages: ChatMessage[],
-    parentEvent?: Event | undefined,
-    streaming?: T,
-  ): Promise<R> {
+  chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<string>> {
+    const { messages, parentEvent, stream } = params;
     //Streaming
-    if (streaming) {
-      if (!this.hasStreaming) {
-        throw Error("No streaming support for this LLM.");
-      }
-      return this.streamChat(messages, parentEvent) as R;
+    if (stream) {
+      return this.streamChat(messages, parentEvent);
     }
 
     //Non-streaming
@@ -748,13 +762,13 @@ export class Anthropic implements LLM {
       message: { content: response.completion.trimStart(), role: "assistant" },
       //^ We're trimming the start because Anthropic often starts with a space in the response
       // That space will be re-added when we generate the next prompt.
-    } as R;
+    };
   }
 
   protected async *streamChat(
     messages: ChatMessage[],
     parentEvent?: Event | undefined,
-  ): AsyncGenerator<string, void, unknown> {
+  ): AsyncIterable<string> {
     // AsyncIterable<AnthropicStreamToken>
     const stream: AsyncIterable<AnthropicStreamToken> =
       await this.session.anthropic.completions.create({
@@ -775,36 +789,9 @@ export class Anthropic implements LLM {
     }
     return;
   }
-
-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    prompt: string,
-    parentEvent?: Event | undefined,
-    streaming?: T,
-  ): Promise<R> {
-    if (streaming) {
-      return this.streamComplete(prompt, parentEvent) as R;
-    }
-    return this.chat(
-      [{ content: prompt, role: "user" }],
-      parentEvent,
-      streaming,
-    ) as R;
-  }
-
-  protected streamComplete(
-    prompt: string,
-    parentEvent?: Event | undefined,
-  ): AsyncGenerator<string, void, unknown> {
-    return this.streamChat([{ content: prompt, role: "user" }], parentEvent);
-  }
 }
 
-export class Portkey implements LLM {
-  hasStreaming: boolean = true;
-
+export class Portkey extends BaseLLM {
   apiKey?: string = undefined;
   baseURL?: string = undefined;
   mode?: string = undefined;
@@ -813,6 +800,7 @@ export class Portkey implements LLM {
   callbackManager?: CallbackManager;
 
   constructor(init?: Partial<Portkey>) {
+    super();
     this.apiKey = init?.apiKey;
     this.baseURL = init?.baseURL;
     this.mode = init?.mode;
@@ -834,45 +822,27 @@ export class Portkey implements LLM {
     throw new Error("metadata not implemented for Portkey");
   }
 
-  async chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    messages: ChatMessage[],
-    parentEvent?: Event | undefined,
-    streaming?: T,
-    params?: Record<string, any>,
-  ): Promise<R> {
-    if (streaming) {
-      return this.streamChat(messages, parentEvent, params) as R;
+  chat(params: LLMChatParamsStreaming): Promise<AsyncIterable<string>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<string>> {
+    const { messages, parentEvent, stream, extraParams } = params;
+    if (stream) {
+      return this.streamChat(messages, parentEvent, extraParams);
     } else {
-      const resolvedParams = params || {};
+      const bodyParams = extraParams || {};
       const response = await this.session.portkey.chatCompletions.create({
         messages,
-        ...resolvedParams,
+        ...bodyParams,
       });
 
       const content = response.choices[0].message?.content ?? "";
       const role = response.choices[0].message?.role || "assistant";
-      return { message: { content, role: role as MessageType } } as R;
+      return { message: { content, role: role as MessageType } };
     }
   }
 
-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    prompt: string,
-    parentEvent?: Event | undefined,
-    streaming?: T,
-  ): Promise<R> {
-    return this.chat(
-      [{ content: prompt, role: "user" }],
-      parentEvent,
-      streaming,
-    );
-  }
-
   async *streamChat(
     messages: ChatMessage[],
     parentEvent?: Event,
@@ -919,11 +889,4 @@ export class Portkey implements LLM {
     }
     return;
   }
-
-  streamComplete(
-    query: string,
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
-    return this.streamChat([{ content: query, role: "user" }], parentEvent);
-  }
 }

packages/core/src/llm/mistral.ts

@@ -41,8 +41,9 @@ export class MistralAISession {
 /**
  * MistralAI LLM implementation
  */
-export class MistralAI implements LLM {
+export class MistralAI implements LLMAI {
   hasStreaming: boolean = true;
+  
 
   // Per completion MistralAI params
   model: keyof typeof ALL_AVAILABLE_MISTRAL_MODELS;
@@ -94,10 +95,7 @@ export class MistralAI implements LLM {
     };
   }
 
-  async chat<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(messages: ChatMessage[], parentEvent?: Event, streaming?: T): Promise<R> {
+  async chat(messages: ChatMessage[], parentEvent?: Event, streaming?: boolean): Promise<ChatResponse> {
     // Streaming
     if (streaming) {
       if (!this.hasStreaming) {
@@ -114,10 +112,7 @@ export class MistralAI implements LLM {
     } as R;
   }
 
-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(prompt: string, parentEvent?: Event, streaming?: T): Promise<R> {
+  async complete(prompt: string, parentEvent?: Event, streaming?: boolean): Promise<ChatResponse> {
     return this.chat(
       [{ content: prompt, role: "user" }],
       parentEvent,

packages/core/src/tests/CallbackManager.test.ts

@@ -8,10 +8,19 @@ import {
 import { OpenAIEmbedding } from "../embeddings";
 import { SummaryIndex } from "../indices/summary";
 import { VectorStoreIndex } from "../indices/vectorStore/VectorStoreIndex";
-import { OpenAI } from "../llm/LLM";
-import { ResponseSynthesizer, SimpleResponseBuilder } from "../synthesizers";
+import { OpenAI } from "../llm/openai";
+
 import { mockEmbeddingModel, mockLlmGeneration } from "./utility/mockOpenAI";
 
+// Mock the OpenAI getOpenAISession function during testing
+jest.mock("../llm/openai", () => {
+  return {
+    getOpenAISession: jest.fn().mockImplementation(() => null),
+  };
+});
+import { ResponseSynthesizer, SimpleResponseBuilder } from "../synthesizers";
+
+
 // Mock the OpenAI getOpenAISession function during testing
 jest.mock("../llm/openai", () => {
   return {