feat: code

Merge remote-tracking branch 'origin/main' into himself65/fix-type
# Conflicts: # packages/core/src/llm/LLM.ts
2026-07-16 07:14:29 -04:00 · 2024-01-15 12:52:25 -06:00 · 2024-01-15 12:49:51 -06:00 · 2024-01-15 12:48:32 -06:00 · 2024-01-15 18:15:44 +07:00 · 2024-01-15 18:06:53 +07:00
46 changed files with 996 additions and 642 deletions
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+fix: openai type might be missing
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+feat: support showing image on chat message
@@ -0,0 +1,5 @@
+---
+"llamaindex": patch
+---
+
+refactor: Updated low-level streaming interface
@@ -35,13 +35,26 @@ const nodesWithScore: NodeWithScore[] = [
  },
 ];

-const response = await responseSynthesizer.synthesize(
-  "What age am I?",
+const response = await responseSynthesizer.synthesize({
+  query: "What age am I?",
  nodesWithScore,
-);
+});
 console.log(response.response);
 ```

+The `synthesize` function also supports streaming, just add `stream: true` as an option:
+
+```typescript
+const stream = await responseSynthesizer.synthesize({
+  query: "What age am I?",
+  nodesWithScore,
+  stream: true,
+});
+for await (const chunk of stream) {
+  process.stdout.write(chunk.response);
+}
+```
+
 ## API Reference

 - [ResponseSynthesizer](../../api/classes/ResponseSynthesizer.md)
@@ -0,0 +1 @@
+label: Observability
@@ -0,0 +1,35 @@
+# Observability
+
+LlamaIndex provides **one-click observability** 🔭 to allow you to build principled LLM applications in a production setting.
+
+A key requirement for principled development of LLM applications over your data (RAG systems, agents) is being able to observe, debug, and evaluate
+your system - both as a whole and for each component.
+
+This feature allows you to seamlessly integrate the LlamaIndex library with powerful observability/evaluation tools offered by our partners.
+Configure a variable once, and you'll be able to do things like the following:
+
+- View LLM/prompt inputs/outputs
+- Ensure that the outputs of any component (LLMs, embeddings) are performing as expected
+- View call traces for both indexing and querying
+
+Each provider has similarities and differences. Take a look below for the full set of guides for each one!
+
+## OpenLLMetry
+
+[OpenLLMetry](https://github.com/traceloop/openllmetry-js) is an open-source project based on OpenTelemetry for tracing and monitoring
+LLM applications. It connects to [all major observability platforms](https://www.traceloop.com/docs/openllmetry/integrations/introduction) and installs in minutes.
+
+### Usage Pattern
+
+```bash
+npm install @traceloop/node-server-sdk
+```
+
+```js
+import * as traceloop from "@traceloop/node-server-sdk";
+
+traceloop.initialize({
+  apiKey: process.env.TRACELOOP_API_KEY,
+  disableBatch: true,
+});
+```
@@ -2,13 +2,15 @@ import { Anthropic } from "llamaindex";

 (async () => {
  const anthropic = new Anthropic();
-  const result = await anthropic.chat([
-    { content: "You want to talk in rhymes.", role: "system" },
-    {
-      content:
-        "How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
-      role: "user",
-    },
-  ]);
+  const result = await anthropic.chat({
+    messages: [
+      { content: "You want to talk in rhymes.", role: "system" },
+      {
+        content:
+          "How much wood would a woodchuck chuck if a woodchuck could chuck wood?",
+        role: "user",
+      },
+    ],
+  });
  console.log(result);
 })();
@@ -25,12 +25,12 @@ import { ChatMessage, LlamaDeuce, OpenAI } from "llamaindex";

  while (true) {
    const next = history.length % 2 === 1 ? gpt4 : l2;
-    const r = await next.chat(
-      history.map(({ content, role }) => ({
+    const r = await next.chat({
+      messages: history.map(({ content, role }) => ({
        content,
        role: next === l2 ? role : role === "user" ? "assistant" : "user",
      })),
-    );
+    });
    history.push({
      content: r.message.content,
      role: next === l2 ? "assistant" : "user",
@@ -21,18 +21,20 @@ async function main() {
    action_items: ["action item 1", "action item 2"],
  };

-  const response = await llm.chat([
-    {
-      role: "system",
-      content: `You are an expert assistant for summarizing and extracting insights from sales call transcripts.\n\nGenerate a valid JSON in the following format:\n\n${JSON.stringify(
-        example,
-      )}`,
-    },
-    {
-      role: "user",
-      content: `Here is the transcript: \n------\n${transcript}\n------`,
-    },
-  ]);
+  const response = await llm.chat({
+    messages: [
+      {
+        role: "system",
+        content: `You are an expert assistant for summarizing and extracting insights from sales call transcripts.\n\nGenerate a valid JSON in the following format:\n\n${JSON.stringify(
+          example,
+        )}`,
+      },
+      {
+        role: "user",
+        content: `Here is the transcript: \n------\n${transcript}\n------`,
+      },
+    ],
+  });

  const json = JSON.parse(response.message.content);

@@ -2,6 +2,8 @@ import { DeuceChatStrategy, LlamaDeuce } from "llamaindex";

 (async () => {
  const deuce = new LlamaDeuce({ chatStrategy: DeuceChatStrategy.META });
-  const result = await deuce.chat([{ content: "Hello, world!", role: "user" }]);
+  const result = await deuce.chat({
+    messages: [{ content: "Hello, world!", role: "user" }],
+  });
  console.log(result);
 })();
@@ -27,9 +27,12 @@ import {
    },
  ];

-  const response = await responseSynthesizer.synthesize(
-    "What age am I?",
+  const stream = await responseSynthesizer.synthesize({
+    query: "What age am I?",
    nodesWithScore,
-  );
-  console.log(response.response);
+    stream: true,
+  });
+  for await (const chunk of stream) {
+    process.stdout.write(chunk.response);
+  }
 })();
@@ -43,19 +43,20 @@ async function rag(llm: LLM, embedModel: BaseEmbedding, query: string) {

  // chat api (non-streaming)
  const llm = new MistralAI({ model: "mistral-tiny" });
-  const response = await llm.chat([
-    { content: "What is the best French cheese?", role: "user" },
-  ]);
+  const response = await llm.chat({
+    messages: [{ content: "What is the best French cheese?", role: "user" }],
+  });
  console.log(response.message.content);

  // chat api (streaming)
-  const stream = await llm.chat(
-    [{ content: "Who is the most renowned French painter?", role: "user" }],
-    undefined,
-    true,
-  );
+  const stream = await llm.chat({
+    messages: [
+      { content: "Who is the most renowned French painter?", role: "user" },
+    ],
+    stream: true,
+  });
  for await (const chunk of stream) {
-    process.stdout.write(chunk);
+    process.stdout.write(chunk.delta);
  }

  // rag
@@ -3,32 +3,34 @@ import { Ollama } from "llamaindex";
 (async () => {
  const llm = new Ollama({ model: "llama2", temperature: 0.75 });
  {
-    const response = await llm.chat([
-      { content: "Tell me a joke.", role: "user" },
-    ]);
+    const response = await llm.chat({
+      messages: [{ content: "Tell me a joke.", role: "user" }],
+    });
    console.log("Response 1:", response.message.content);
  }
  {
-    const response = await llm.complete("How are you?");
-    console.log("Response 2:", response.message.content);
+    const response = await llm.complete({ prompt: "How are you?" });
+    console.log("Response 2:", response.text);
  }
  {
-    const response = await llm.chat(
-      [{ content: "Tell me a joke.", role: "user" }],
-      undefined,
-      true,
-    );
+    const response = await llm.chat({
+      messages: [{ content: "Tell me a joke.", role: "user" }],
+      stream: true,
+    });
    console.log("Response 3:");
    for await (const message of response) {
-      process.stdout.write(message); // no newline
+      process.stdout.write(message.delta); // no newline
    }
    console.log(); // newline
  }
  {
-    const response = await llm.complete("How are you?", undefined, true);
+    const response = await llm.complete({
+      prompt: "How are you?",
+      stream: true,
+    });
    console.log("Response 4:");
    for await (const message of response) {
-      process.stdout.write(message); // no newline
+      process.stdout.write(message.text); // no newline
    }
    console.log(); // newline
  }
@@ -4,12 +4,12 @@ import { OpenAI } from "llamaindex";
  const llm = new OpenAI({ model: "gpt-4-1106-preview", temperature: 0.1 });

  // complete api
-  const response1 = await llm.complete("How are you?");
-  console.log(response1.message.content);
+  const response1 = await llm.complete({ prompt: "How are you?" });
+  console.log(response1.text);

  // chat api
-  const response2 = await llm.chat([
-    { content: "Tell me a joke.", role: "user" },
-  ]);
+  const response2 = await llm.chat({
+    messages: [{ content: "Tell me a joke.", role: "user" }],
+  });
  console.log(response2.message.content);
 })();
@@ -12,11 +12,14 @@ import { Portkey } from "llamaindex";
      },
    ],
  });
-  const result = portkey.streamChat([
-    { role: "system", content: "You are a helpful assistant." },
-    { role: "user", content: "Tell me a joke." },
-  ]);
+  const result = await portkey.chat({
+    messages: [
+      { role: "system", content: "You are a helpful assistant." },
+      { role: "user", content: "Tell me a joke." },
+    ],
+    stream: true,
+  });
  for await (const res of result) {
-    process.stdout.write(res);
+    process.stdout.write(res.delta);
  }
 })();
@@ -6,8 +6,8 @@ const together = new TogetherLLM({
 });

 (async () => {
-  const generator = await together.chat(
-    [
+  const generator = await together.chat({
+    messages: [
      {
        role: "system",
        content: "You are an AI assistant",
@@ -17,12 +17,11 @@ const together = new TogetherLLM({
        content: "Tell me about San Francisco",
      },
    ],
-    undefined,
-    true,
-  );
+    stream: true,
+  });
  console.log("Chatting with Together AI...");
  for await (const message of generator) {
-    process.stdout.write(message);
+    process.stdout.write(message.delta);
  }
  const embedding = new TogetherEmbedding();
  const vector = await embedding.getTextEmbedding("Hello world!");
@@ -4,12 +4,12 @@ import { OpenAI } from "llamaindex";
  const llm = new OpenAI({ model: "gpt-4-vision-preview", temperature: 0.1 });

  // complete api
-  const response1 = await llm.complete("How are you?");
-  console.log(response1.message.content);
+  const response1 = await llm.complete({ prompt: "How are you?" });
+  console.log(response1.text);

  // chat api
-  const response2 = await llm.chat([
-    { content: "Tell me a joke!", role: "user" },
-  ]);
+  const response2 = await llm.chat({
+    messages: [{ content: "Tell me a joke!", role: "user" }],
+  });
  console.log(response2.message.content);
 })();
@@ -69,7 +69,7 @@ export class SimpleChatEngine implements ChatEngine {
    //Non-streaming option
    chatHistory = chatHistory ?? this.chatHistory;
    chatHistory.push({ content: message, role: "user" });
-    const response = await this.llm.chat(chatHistory, undefined);
+    const response = await this.llm.chat({ messages: chatHistory });
    chatHistory.push(response.message);
    this.chatHistory = chatHistory;
    return new Response(response.message.content) as R;
@@ -81,16 +81,15 @@ export class SimpleChatEngine implements ChatEngine {
  ): AsyncGenerator<string, void, unknown> {
    chatHistory = chatHistory ?? this.chatHistory;
    chatHistory.push({ content: message, role: "user" });
-    const response_generator = await this.llm.chat(
-      chatHistory,
-      undefined,
-      true,
-    );
+    const response_generator = await this.llm.chat({
+      messages: chatHistory,
+      stream: true,
+    });

    var accumulator: string = "";
    for await (const part of response_generator) {
-      accumulator += part;
-      yield part;
+      accumulator += part.delta;
+      yield part.delta;
    }

    chatHistory.push({ content: accumulator, role: "assistant" });
@@ -136,12 +135,12 @@ export class CondenseQuestionChatEngine implements ChatEngine {
  private async condenseQuestion(chatHistory: ChatMessage[], question: string) {
    const chatHistoryStr = messagesToHistoryStr(chatHistory);

-    return this.serviceContext.llm.complete(
-      defaultCondenseQuestionPrompt({
+    return this.serviceContext.llm.complete({
+      prompt: defaultCondenseQuestionPrompt({
        question: question,
        chatHistory: chatHistoryStr,
      }),
-    );
+    });
  }

  async chat<
@@ -156,7 +155,7 @@ export class CondenseQuestionChatEngine implements ChatEngine {

    const condensedQuestion = (
      await this.condenseQuestion(chatHistory, extractText(message))
-    ).message.content;
+    ).text;

    const response = await this.queryEngine.query(condensedQuestion);

@@ -283,10 +282,10 @@ export class ContextChatEngine implements ChatEngine {

    chatHistory.push({ content: message, role: "user" });

-    const response = await this.chatModel.chat(
-      [context.message, ...chatHistory],
+    const response = await this.chatModel.chat({
+      messages: [context.message, ...chatHistory],
      parentEvent,
-    );
+    });
    chatHistory.push(response.message);

    this.chatHistory = chatHistory;
@@ -315,15 +314,15 @@ export class ContextChatEngine implements ChatEngine {

    chatHistory.push({ content: message, role: "user" });

-    const response_stream = await this.chatModel.chat(
-      [context.message, ...chatHistory],
+    const response_stream = await this.chatModel.chat({
+      messages: [context.message, ...chatHistory],
      parentEvent,
-      true,
-    );
+      stream: true,
+    });
    var accumulator: string = "";
    for await (const part of response_stream) {
-      accumulator += part;
-      yield part;
+      accumulator += part.delta;
+      yield part.delta;
    }

    chatHistory.push({ content: accumulator, role: "assistant" });
@@ -399,7 +398,7 @@ export class HistoryChatEngine {
      message,
      chatHistory,
    );
-    const response = await this.llm.chat(requestMessages);
+    const response = await this.llm.chat({ messages: requestMessages });
    chatHistory.addMessage(response.message);
    return new Response(response.message.content) as R;
  }
@@ -412,16 +411,15 @@ export class HistoryChatEngine {
      message,
      chatHistory,
    );
-    const response_stream = await this.llm.chat(
-      requestMessages,
-      undefined,
-      true,
-    );
+    const response_stream = await this.llm.chat({
+      messages: requestMessages,
+      stream: true,
+    });

    var accumulator = "";
    for await (const part of response_stream) {
-      accumulator += part;
-      yield part;
+      accumulator += part.delta;
+      yield part.delta;
    }
    chatHistory.addMessage({
      content: accumulator,
@@ -99,7 +99,7 @@ export class SummaryChatHistory implements ChatHistory {
      messagesToSummarize.shift();
    } while (this.llm.tokens(promptMessages) > this.tokensToSummarize);

-    const response = await this.llm.chat(promptMessages);
+    const response = await this.llm.chat({ messages: promptMessages });
    return { content: response.message.content, role: "memory" };
  }

@@ -76,8 +76,12 @@ export class RetrieverQueryEngine implements BaseQueryEngine {
      type: "wrapper",
      tags: ["final"],
    };
-    const nodes = await this.retrieve(query, _parentEvent);
-    return this.responseSynthesizer.synthesize(query, nodes, _parentEvent);
+    const nodesWithScore = await this.retrieve(query, _parentEvent);
+    return this.responseSynthesizer.synthesize({
+      query,
+      nodesWithScore,
+      parentEvent: _parentEvent,
+    });
  }
 }

@@ -153,10 +157,14 @@ export class SubQuestionQueryEngine implements BaseQueryEngine {
      subQuestions.map((subQ) => this.querySubQ(subQ, subQueryParentEvent)),
    );

-    const nodes = subQNodes
+    const nodesWithScore = subQNodes
      .filter((node) => node !== null)
      .map((node) => node as NodeWithScore);
-    return this.responseSynthesizer.synthesize(query, nodes, parentEvent);
+    return this.responseSynthesizer.synthesize({
+      query,
+      nodesWithScore,
+      parentEvent,
+    });
  }

  private async querySubQ(
@@ -41,13 +41,13 @@ export class LLMQuestionGenerator implements BaseQuestionGenerator {
    const toolsStr = buildToolsText(tools);
    const queryStr = query;
    const prediction = (
-      await this.llm.complete(
-        this.prompt({
+      await this.llm.complete({
+        prompt: this.prompt({
          toolsStr,
          queryStr,
        }),
-      )
-    ).message.content;
+      })
+    ).text;

    const structuredOutput = this.outputParser.parse(prediction);

@@ -1,7 +1,7 @@
 import { BaseNode } from "./Node";

 /**
- * Respone is the output of a LLM
+ * Response is the output of a LLM
 */
 export class Response {
  response: string;
@@ -13,8 +13,8 @@ export enum OpenAIEmbeddingModelType {
  TEXT_EMBED_ADA_002 = "text-embedding-ada-002",
 }

-export class OpenAIEmbedding extends BaseEmbedding {
-  model: OpenAIEmbeddingModelType | string;
+export abstract class OpenAIEmbeddingLike extends BaseEmbedding {
+  abstract model: string;

  // OpenAI session params
  apiKey?: string = undefined;
@@ -27,15 +27,47 @@ export class OpenAIEmbedding extends BaseEmbedding {

  session: OpenAISession;

-  constructor(init?: Partial<OpenAIEmbedding> & { azure?: AzureOpenAIConfig }) {
+  constructor(init?: Partial<OpenAIEmbeddingLike>) {
    super();

-    this.model = OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002;
-
    this.maxRetries = init?.maxRetries ?? 10;
    this.timeout = init?.timeout ?? 60 * 1000; // Default is 60 seconds
    this.additionalSessionOptions = init?.additionalSessionOptions;

+    this.apiKey = init?.apiKey ?? undefined;
+    this.session =
+      init?.session ??
+      getOpenAISession({
+        apiKey: this.apiKey,
+        maxRetries: this.maxRetries,
+        timeout: this.timeout,
+        ...this.additionalSessionOptions,
+      });
+  }
+
+  private async getOpenAIEmbedding(input: string) {
+    const { data } = await this.session.openai.embeddings.create({
+      model: this.model,
+      input,
+    });
+
+    return data[0].embedding;
+  }
+
+  async getTextEmbedding(text: string): Promise<number[]> {
+    return this.getOpenAIEmbedding(text);
+  }
+
+  async getQueryEmbedding(query: string): Promise<number[]> {
+    return this.getOpenAIEmbedding(query);
+  }
+}
+
+export class OpenAIEmbedding extends OpenAIEmbeddingLike {
+  public override model: OpenAIEmbeddingModelType;
+  constructor(init?: Partial<OpenAIEmbedding> & { azure?: AzureOpenAIConfig }) {
+    super(init);
+    this.model = init?.model ?? OpenAIEmbeddingModelType.TEXT_EMBED_ADA_002;
    if (init?.azure || shouldUseAzure()) {
      const azureConfig = getAzureConfigFromEnv({
        ...init?.azure,
@@ -60,33 +92,6 @@ export class OpenAIEmbedding extends BaseEmbedding {
          defaultQuery: { "api-version": azureConfig.apiVersion },
          ...this.additionalSessionOptions,
        });
-    } else {
-      this.apiKey = init?.apiKey ?? undefined;
-      this.session =
-        init?.session ??
-        getOpenAISession({
-          apiKey: this.apiKey,
-          maxRetries: this.maxRetries,
-          timeout: this.timeout,
-          ...this.additionalSessionOptions,
-        });
    }
  }
-
-  private async getOpenAIEmbedding(input: string) {
-    const { data } = await this.session.openai.embeddings.create({
-      model: this.model,
-      input,
-    });
-
-    return data[0].embedding;
-  }
-
-  async getTextEmbedding(text: string): Promise<number[]> {
-    return this.getOpenAIEmbedding(text);
-  }
-
-  async getQueryEmbedding(query: string): Promise<number[]> {
-    return this.getOpenAIEmbedding(query);
-  }
 }
@@ -1,8 +1,8 @@
-import { OpenAIEmbedding } from "./OpenAIEmbedding";
+import { OpenAIEmbeddingLike } from "./OpenAIEmbedding";

-export class TogetherEmbedding extends OpenAIEmbedding {
+export class TogetherEmbedding extends OpenAIEmbeddingLike {
  override model: string;
-  constructor(init?: Partial<OpenAIEmbedding>) {
+  constructor(init?: Partial<TogetherEmbedding>) {
    super({
      apiKey: process.env.TOGETHER_API_KEY,
      ...init,
@@ -149,12 +149,12 @@ export class KeywordTableIndex extends BaseIndex<KeywordTable> {
    text: string,
    serviceContext: ServiceContext,
  ): Promise<Set<string>> {
-    const response = await serviceContext.llm.complete(
-      defaultKeywordExtractPrompt({
+    const response = await serviceContext.llm.complete({
+      prompt: defaultKeywordExtractPrompt({
        context: text,
      }),
-    );
-    return extractKeywordsGivenResponse(response.message.content, "KEYWORDS:");
+    });
+    return extractKeywordsGivenResponse(response.text, "KEYWORDS:");
  }

  /**
@@ -86,16 +86,13 @@ abstract class BaseKeywordTableRetriever implements BaseRetriever {
 // Extracts keywords using LLMs.
 export class KeywordTableLLMRetriever extends BaseKeywordTableRetriever {
  async getKeywords(query: string): Promise<string[]> {
-    const response = await this.serviceContext.llm.complete(
-      this.queryKeywordExtractTemplate({
+    const response = await this.serviceContext.llm.complete({
+      prompt: this.queryKeywordExtractTemplate({
        question: query,
        maxKeywords: this.maxKeywordsPerQuery,
      }),
-    );
-    const keywords = extractKeywordsGivenResponse(
-      response.message.content,
-      "KEYWORDS:",
-    );
+    });
+    const keywords = extractKeywordsGivenResponse(response.text, "KEYWORDS:");
    return [...keywords];
  }
 }
@@ -89,8 +89,10 @@ export class SummaryIndexLLMRetriever implements BaseRetriever {
      const fmtBatchStr = this.formatNodeBatchFn(nodesBatch);
      const input = { context: fmtBatchStr, query: query };
      const rawResponse = (
-        await this.serviceContext.llm.complete(this.choiceSelectPrompt(input))
-      ).message.content;
+        await this.serviceContext.llm.complete({
+          prompt: this.choiceSelectPrompt(input),
+        })
+      ).text;

      // parseResult is a map from doc number to relevance score
      const parseResult = this.parseChoiceSelectAnswerFn(
@@ -10,8 +10,7 @@ import {

 import { ChatCompletionMessageParam } from "openai/resources";
 import { LLMOptions } from "portkey-ai";
-import { MessageContent } from "../ChatEngine";
-import { globalsHelper, Tokenizers } from "../GlobalsHelper";
+import { Tokenizers, globalsHelper } from "../GlobalsHelper";
 import {
  ANTHROPIC_AI_PROMPT,
  ANTHROPIC_HUMAN_PROMPT,
@@ -25,8 +24,8 @@ import {
  getAzureModel,
  shouldUseAzure,
 } from "./azure";
-import { getOpenAISession, OpenAISession } from "./openai";
-import { getPortkeySession, PortkeySession } from "./portkey";
+import { OpenAISession, getOpenAISession } from "./openai";
+import { PortkeySession, getPortkeySession } from "./portkey";
 import { ReplicateSession } from "./replicate";

 export type MessageType =
@@ -45,11 +44,16 @@ export interface ChatMessage {
 export interface ChatResponse {
  message: ChatMessage;
  raw?: Record<string, any>;
-  delta?: string;
 }

-// NOTE in case we need CompletionResponse to diverge from ChatResponse in the future
-export type CompletionResponse = ChatResponse;
+export interface ChatResponseChunk {
+  delta: string;
+}
+
+export interface CompletionResponse {
+  text: string;
+  raw?: Record<string, any>;
+}

 export interface LLMMetadata {
  model: string;
@@ -60,40 +64,59 @@ 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: any;
+  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<ChatResponseChunk>>;
+  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<CompletionResponse>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;

  /**
   * Calculates the number of tokens needed for the given chat messages
@@ -101,6 +124,50 @@ export interface LLM {
  tokens(messages: ChatMessage[]): number;
 }

+export abstract class BaseLLM implements LLM {
+  abstract metadata: LLMMetadata;
+
+  private async *chatToComplete(
+    stream: AsyncIterable<ChatResponseChunk>,
+  ): AsyncIterable<CompletionResponse> {
+    for await (const chunk of stream) {
+      yield { text: chunk.delta };
+    }
+  }
+
+  complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
+    const { prompt, parentEvent, stream } = params;
+    if (stream) {
+      const stream = await this.chat({
+        messages: [{ content: prompt, role: "user" }],
+        parentEvent,
+        stream: true,
+      });
+      return this.chatToComplete(stream);
+    }
+    const chatResponse = await this.chat({
+      messages: [{ content: prompt, role: "user" }],
+      parentEvent,
+    });
+    return { text: chatResponse.message.content as string };
+  }
+
+  abstract chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  abstract chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+
+  abstract tokens(messages: ChatMessage[]): number;
+}
+
 export const GPT4_MODELS = {
  "gpt-4": { contextWindow: 8192 },
  "gpt-4-32k": { contextWindow: 32768 },
@@ -117,25 +184,28 @@ export const GPT35_MODELS = {
 /**
 * We currently support GPT-3.5 and GPT-4 models
 */
-export const ALL_AVAILABLE_OPENAI_MODELS = {
+export const ALL_AVAILABLE_OPENAI_MODELS: Record<
+  string,
+  {
+    contextWindow: number;
+  }
+> = {
  ...GPT4_MODELS,
  ...GPT35_MODELS,
 };

-/**
- * OpenAI LLM implementation
- */
-export class OpenAI implements LLM {
+type OpenAIModel = keyof typeof GPT4_MODELS | keyof typeof GPT35_MODELS;
+
+export abstract class OpenAILike extends BaseLLM implements LLM {
  hasStreaming: boolean = true;

-  // Per completion OpenAI params
-  model: keyof typeof ALL_AVAILABLE_OPENAI_MODELS | string;
+  abstract model: string;
  temperature: number;
  topP: number;
  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
@@ -150,12 +220,8 @@ export class OpenAI implements LLM {

  callbackManager?: CallbackManager;

-  constructor(
-    init?: Partial<OpenAI> & {
-      azure?: AzureOpenAIConfig;
-    },
-  ) {
-    this.model = init?.model ?? "gpt-3.5-turbo";
+  constructor(init?: Partial<OpenAILike>) {
+    super();
    this.temperature = init?.temperature ?? 0.1;
    this.topP = init?.topP ?? 1;
    this.maxTokens = init?.maxTokens ?? undefined;
@@ -165,56 +231,26 @@ export class OpenAI implements LLM {
    this.additionalChatOptions = init?.additionalChatOptions;
    this.additionalSessionOptions = init?.additionalSessionOptions;

-    if (init?.azure || shouldUseAzure()) {
-      const azureConfig = getAzureConfigFromEnv({
-        ...init?.azure,
-        model: getAzureModel(this.model),
+    this.apiKey = init?.apiKey ?? undefined;
+    this.session =
+      init?.session ??
+      getOpenAISession({
+        apiKey: this.apiKey,
+        maxRetries: this.maxRetries,
+        timeout: this.timeout,
+        ...this.additionalSessionOptions,
      });

-      if (!azureConfig.apiKey) {
-        throw new Error(
-          "Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
-        );
-      }
-
-      this.apiKey = azureConfig.apiKey;
-      this.session =
-        init?.session ??
-        getOpenAISession({
-          azure: true,
-          apiKey: this.apiKey,
-          baseURL: getAzureBaseUrl(azureConfig),
-          maxRetries: this.maxRetries,
-          timeout: this.timeout,
-          defaultQuery: { "api-version": azureConfig.apiVersion },
-          ...this.additionalSessionOptions,
-        });
-    } else {
-      this.apiKey = init?.apiKey ?? undefined;
-      this.session =
-        init?.session ??
-        getOpenAISession({
-          apiKey: this.apiKey,
-          maxRetries: this.maxRetries,
-          timeout: this.timeout,
-          ...this.additionalSessionOptions,
-        });
-    }
-
    this.callbackManager = init?.callbackManager;
  }

  get metadata() {
-    const contextWindow =
-      ALL_AVAILABLE_OPENAI_MODELS[
-        this.model as keyof typeof ALL_AVAILABLE_OPENAI_MODELS
-      ]?.contextWindow ?? 1024;
    return {
      model: this.model,
      temperature: this.temperature,
      topP: this.topP,
      maxTokens: this.maxTokens,
-      contextWindow,
+      contextWindow: ALL_AVAILABLE_OPENAI_MODELS[this.model].contextWindow,
      tokenizer: Tokenizers.CL100K_BASE,
    };
  }
@@ -251,10 +287,14 @@ 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<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, parentEvent, stream } = params;
    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
      model: this.model,
      temperature: this.temperature,
@@ -270,11 +310,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({
@@ -285,27 +322,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<ChatResponseChunk> {
    const baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams = {
      model: this.model,
      temperature: this.temperature,
@@ -360,17 +383,48 @@ export class OpenAI implements LLM {

      idx_counter++;

-      yield part.choices[0].delta.content ? part.choices[0].delta.content : "";
+      yield {
+        delta: part.choices[0].delta.content ?? "",
+      };
    }
    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 class OpenAI extends OpenAILike {
+  model: OpenAIModel;
+  constructor(
+    init?: Partial<OpenAI> & {
+      azure?: AzureOpenAIConfig;
+    },
+  ) {
+    super(init);
+    this.model = init?.model ?? "gpt-3.5-turbo";
+    if (init?.azure || shouldUseAzure()) {
+      const azureConfig = getAzureConfigFromEnv({
+        ...init?.azure,
+        model: getAzureModel(this.model),
+      });
+
+      if (!azureConfig.apiKey) {
+        throw new Error(
+          "Azure API key is required for OpenAI Azure models. Please set the AZURE_OPENAI_KEY environment variable.",
+        );
+      }
+
+      this.apiKey = azureConfig.apiKey;
+      this.session =
+        init?.session ??
+        getOpenAISession({
+          azure: true,
+          apiKey: this.apiKey,
+          baseURL: getAzureBaseUrl(azureConfig),
+          maxRetries: this.maxRetries,
+          timeout: this.timeout,
+          defaultQuery: { "api-version": azureConfig.apiVersion },
+          ...this.additionalSessionOptions,
+        });
+    }
  }
 }

@@ -429,16 +483,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 ??
@@ -451,7 +505,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 {
@@ -596,10 +649,14 @@ 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<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, parentEvent, stream } = params;
    const api = ALL_AVAILABLE_LLAMADEUCE_MODELS[this.model]
      .replicateApi as `${string}/${string}:${string}`;

@@ -621,6 +678,9 @@ If a question does not make any sense, or is not factually coherent, explain why
    }

    //TODO: Add streaming for this
+    if (stream) {
+      throw new Error("Streaming not supported for LlamaDeuce");
+    }

    //Non-streaming
    const response = await this.replicateSession.replicate.run(
@@ -633,14 +693,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);
+    };
  }
 }

@@ -654,9 +707,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;
@@ -672,6 +723,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
@@ -723,20 +775,17 @@ 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<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    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
@@ -752,13 +801,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<ChatResponseChunk> {
    // AsyncIterable<AnthropicStreamToken>
    const stream: AsyncIterable<AnthropicStreamToken> =
      await this.session.anthropic.completions.create({
@@ -775,40 +824,13 @@ export class Anthropic implements LLM {
      //TODO: LLM Stream Callback, pending re-work.

      idx_counter++;
-      yield part.completion;
+      yield { delta: part.completion };
    }
    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;
@@ -817,6 +839,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;
@@ -838,50 +861,34 @@ 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<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    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,
    params?: Record<string, any>,
-  ): AsyncGenerator<string, void, unknown> {
+  ): AsyncIterable<ChatResponseChunk> {
    // Wrapping the stream in a callback.
    const onLLMStream = this.callbackManager?.onLLMStream
      ? this.callbackManager.onLLMStream
@@ -919,15 +926,8 @@ export class Portkey implements LLM {

      idx_counter++;

-      yield part.choices[0].delta?.content ?? "";
+      yield { delta: part.choices[0].delta?.content ?? "" };
    }
    return;
  }
-
-  streamComplete(
-    query: string,
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
-    return this.streamChat([{ content: query, role: "user" }], parentEvent);
-  }
 }
@@ -4,7 +4,14 @@ import {
  EventType,
  StreamCallbackResponse,
 } from "../callbacks/CallbackManager";
-import { ChatMessage, ChatResponse, LLM } from "./LLM";
+import {
+  BaseLLM,
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+} from "./LLM";

 export const ALL_AVAILABLE_MISTRAL_MODELS = {
  "mistral-tiny": { contextWindow: 32000 },
@@ -41,9 +48,7 @@ export class MistralAISession {
 /**
 * MistralAI LLM implementation
 */
-export class MistralAI implements LLM {
-  hasStreaming: boolean = true;
-
+export class MistralAI extends BaseLLM {
  // Per completion MistralAI params
  model: keyof typeof ALL_AVAILABLE_MISTRAL_MODELS;
  temperature: number;
@@ -57,6 +62,7 @@ export class MistralAI implements LLM {
  private session: MistralAISession;

  constructor(init?: Partial<MistralAI>) {
+    super();
    this.model = init?.model ?? "mistral-small";
    this.temperature = init?.temperature ?? 0.1;
    this.topP = init?.topP ?? 1;
@@ -94,16 +100,17 @@ 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> {
+  chat(
+    params: LLMChatParamsStreaming,
+  ): Promise<AsyncIterable<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, 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(params);
    }
    // Non-streaming
    const client = await this.session.getClient();
@@ -111,24 +118,13 @@ export class MistralAI implements LLM {
    const message = response.choices[0].message;
    return {
      message,
-    } 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,
-    );
-  }
-
-  protected async *streamChat(
-    messages: ChatMessage[],
-    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
+  protected async *streamChat({
+    messages,
+    parentEvent,
+  }: LLMChatParamsStreaming): AsyncIterable<ChatResponseChunk> {
    //Now let's wrap our stream in a callback
    const onLLMStream = this.callbackManager?.onLLMStream
      ? this.callbackManager.onLLMStream
@@ -163,16 +159,10 @@ export class MistralAI implements LLM {

      idx_counter++;

-      yield part.choices[0].delta.content ?? "";
+      yield {
+        delta: part.choices[0].delta.content ?? "",
+      };
    }
    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);
-  }
 }
@@ -1,11 +1,27 @@
 import { ok } from "node:assert";
-import { MessageContent } from "../ChatEngine";
 import { CallbackManager, Event } from "../callbacks/CallbackManager";
 import { BaseEmbedding } from "../embeddings";
-import { ChatMessage, ChatResponse, LLM, LLMMetadata } from "./LLM";
+import {
+  ChatMessage,
+  ChatResponse,
+  ChatResponseChunk,
+  CompletionResponse,
+  LLM,
+  LLMChatParamsNonStreaming,
+  LLMChatParamsStreaming,
+  LLMCompletionParamsNonStreaming,
+  LLMCompletionParamsStreaming,
+  LLMMetadata,
+} from "./LLM";

-const messageAccessor = (data: any) => data.message.content;
-const completionAccessor = (data: any) => data.response;
+const messageAccessor = (data: any): ChatResponseChunk => {
+  return {
+    delta: data.message.content,
+  };
+};
+const completionAccessor = (data: any): CompletionResponse => {
+  return { text: data.response };
+};

 // https://github.com/jmorganca/ollama
 export class Ollama extends BaseEmbedding implements LLM {
@@ -43,21 +59,21 @@ export class Ollama extends BaseEmbedding 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<ChatResponseChunk>>;
+  chat(params: LLMChatParamsNonStreaming): Promise<ChatResponse>;
+  async chat(
+    params: LLMChatParamsNonStreaming | LLMChatParamsStreaming,
+  ): Promise<ChatResponse | AsyncIterable<ChatResponseChunk>> {
+    const { messages, parentEvent, stream } = params;
    const payload = {
      model: this.model,
      messages: messages.map((message) => ({
        role: message.role,
        content: message.content,
      })),
-      stream: !!streaming,
+      stream: !!stream,
      options: {
        temperature: this.temperature,
        num_ctx: this.contextWindow,
@@ -73,7 +89,7 @@ export class Ollama extends BaseEmbedding implements LLM {
        "Content-Type": "application/json",
      },
    });
-    if (!streaming) {
+    if (!stream) {
      const raw = await response.json();
      const { message } = raw;
      return {
@@ -82,20 +98,20 @@ export class Ollama extends BaseEmbedding implements LLM {
          content: message.content,
        },
        raw,
-      } satisfies ChatResponse as R;
+      };
    } else {
      const stream = response.body;
      ok(stream, "stream is null");
      ok(stream instanceof ReadableStream, "stream is not readable");
-      return this.streamChat(stream, messageAccessor, parentEvent) as R;
+      return this.streamChat(stream, messageAccessor, parentEvent);
    }
  }

-  private async *streamChat(
+  private async *streamChat<T>(
    stream: ReadableStream<Uint8Array>,
-    accessor: (data: any) => string,
+    accessor: (data: any) => T,
    parentEvent?: Event,
-  ): AsyncGenerator<string, void, unknown> {
+  ): AsyncIterable<T> {
    const reader = stream.getReader();
    while (true) {
      const { done, value } = await reader.read();
@@ -119,18 +135,20 @@ export class Ollama extends BaseEmbedding implements LLM {
    }
  }

-  async complete<
-    T extends boolean | undefined = undefined,
-    R = T extends true ? AsyncGenerator<string, void, unknown> : ChatResponse,
-  >(
-    prompt: MessageContent,
-    parentEvent?: Event | undefined,
-    streaming?: T | undefined,
-  ): Promise<R> {
+  complete(
+    params: LLMCompletionParamsStreaming,
+  ): Promise<AsyncIterable<CompletionResponse>>;
+  complete(
+    params: LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse>;
+  async complete(
+    params: LLMCompletionParamsStreaming | LLMCompletionParamsNonStreaming,
+  ): Promise<CompletionResponse | AsyncIterable<CompletionResponse>> {
+    const { prompt, parentEvent, stream } = params;
    const payload = {
      model: this.model,
      prompt: prompt,
-      stream: !!streaming,
+      stream: !!stream,
      options: {
        temperature: this.temperature,
        num_ctx: this.contextWindow,
@@ -146,20 +164,17 @@ export class Ollama extends BaseEmbedding implements LLM {
        "Content-Type": "application/json",
      },
    });
-    if (!streaming) {
+    if (!stream) {
      const raw = await response.json();
      return {
-        message: {
-          role: "assistant",
-          content: raw.response,
-        },
+        text: raw.response,
        raw,
-      } satisfies ChatResponse as R;
+      };
    } else {
      const stream = response.body;
      ok(stream, "stream is null");
      ok(stream instanceof ReadableStream, "stream is not readable");
-      return this.streamChat(stream, completionAccessor, parentEvent) as R;
+      return this.streamChat(stream, completionAccessor, parentEvent);
    }
  }

@@ -1,4 +1,3 @@
-import _ from "lodash";
 import OpenAI, { ClientOptions } from "openai";

 export class AzureOpenAI extends OpenAI {
@@ -35,8 +34,10 @@ export class OpenAISession {
 // I'm not 100% sure this is necessary vs. just starting a new session
 // every time we make a call. They say they try to reuse connections
 // so in theory this is more efficient, but we should test it in the future.
-let defaultOpenAISession: { session: OpenAISession; options: ClientOptions }[] =
-  [];
+let defaultOpenAISession: {
+  session: OpenAISession;
+  options: ClientOptions;
+} | null = null;

 /**
 * Get a session for the OpenAI API. If one already exists with the same options,
@@ -47,14 +48,10 @@ let defaultOpenAISession: { session: OpenAISession; options: ClientOptions }[] =
 export function getOpenAISession(
  options: ClientOptions & { azure?: boolean } = {},
 ) {
-  let session = defaultOpenAISession.find((session) => {
-    return _.isEqual(session.options, options);
-  })?.session;
-
-  if (!session) {
-    session = new OpenAISession(options);
-    defaultOpenAISession.push({ session, options });
+  if (!defaultOpenAISession) {
+    const session = new OpenAISession(options);
+    defaultOpenAISession = { session, options };
  }

-  return session;
+  return defaultOpenAISession.session;
 }
@@ -1,7 +1,13 @@
-import { OpenAI } from "./LLM";
+import { Tokenizers } from "../GlobalsHelper";
+import { OpenAILike } from "./LLM";

-export class TogetherLLM extends OpenAI {
-  constructor(init?: Partial<OpenAI>) {
+export class TogetherLLM extends OpenAILike {
+  override model: string;
+  constructor(
+    init?: Partial<TogetherLLM> & {
+      model?: string;
+    },
+  ) {
    super({
      ...init,
      apiKey: process.env.TOGETHER_API_KEY,
@@ -10,5 +16,18 @@ export class TogetherLLM extends OpenAI {
        baseURL: "https://api.together.xyz/v1",
      },
    });
+    this.model = init?.model ?? '"mistralai/Mixtral-8x7B-Instruct-v0.1';
+  }
+
+  get metadata() {
+    return {
+      model: this.model,
+      temperature: this.temperature,
+      topP: this.topP,
+      maxTokens: this.maxTokens,
+      // todo: cannot find context window in documentation
+      contextWindow: 1024,
+      tokenizer: Tokenizers.CL100K_BASE,
+    };
  }
 }
@@ -0,0 +1,10 @@
+// TODO: use for LLM.ts
+
+export async function* streamConverter<S, D>(
+  stream: AsyncIterable<S>,
+  converter: (s: S) => D,
+): AsyncIterable<D> {
+  for await (const data of stream) {
+    yield converter(data);
+  }
+}
@@ -1,16 +1,14 @@
 import { MessageContentDetail } from "../ChatEngine";
-import {
-  ImageNode,
-  MetadataMode,
-  NodeWithScore,
-  splitNodesByType,
-} from "../Node";
+import { ImageNode, MetadataMode, splitNodesByType } from "../Node";
 import { Response } from "../Response";
 import { ServiceContext, serviceContextFromDefaults } from "../ServiceContext";
-import { Event } from "../callbacks/CallbackManager";
 import { imageToDataUrl } from "../embeddings";
 import { TextQaPrompt, defaultTextQaPrompt } from "./../Prompt";
-import { BaseSynthesizer } from "./types";
+import {
+  BaseSynthesizer,
+  SynthesizeParamsNonStreaming,
+  SynthesizeParamsStreaming,
+} from "./types";

 export class MultiModalResponseSynthesizer implements BaseSynthesizer {
  serviceContext: ServiceContext;
@@ -27,11 +25,21 @@ export class MultiModalResponseSynthesizer implements BaseSynthesizer {
    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
  }

-  async synthesize(
-    query: string,
-    nodesWithScore: NodeWithScore[],
-    parentEvent?: Event,
-  ): Promise<Response> {
+  synthesize(
+    params: SynthesizeParamsStreaming,
+  ): Promise<AsyncIterable<Response>>;
+  synthesize(params: SynthesizeParamsNonStreaming): Promise<Response>;
+  async synthesize({
+    query,
+    nodesWithScore,
+    parentEvent,
+    stream,
+  }: SynthesizeParamsStreaming | SynthesizeParamsNonStreaming): Promise<
+    AsyncIterable<Response> | Response
+  > {
+    if (stream) {
+      throw new Error("streaming not implemented");
+    }
    const nodes = nodesWithScore.map(({ node }) => node);
    const { imageNodes, textNodes } = splitNodesByType(nodes);
    const textChunks = textNodes.map((node) =>
@@ -54,7 +62,10 @@ export class MultiModalResponseSynthesizer implements BaseSynthesizer {
      { type: "text", text: textPrompt },
      ...images,
    ];
-    let response = await this.serviceContext.llm.complete(prompt, parentEvent);
-    return new Response(response.message.content, nodes);
+    let response = await this.serviceContext.llm.complete({
+      prompt,
+      parentEvent,
+    });
+    return new Response(response.text, nodes);
  }
 }
@@ -1,15 +1,20 @@
-import { Event } from "../callbacks/CallbackManager";
-import { MetadataMode, NodeWithScore } from "../Node";
+import { MetadataMode } from "../Node";
 import { Response } from "../Response";
 import { ServiceContext, serviceContextFromDefaults } from "../ServiceContext";
-import { BaseResponseBuilder, getResponseBuilder } from "./builders";
-import { BaseSynthesizer } from "./types";
+import { streamConverter } from "../llm/utils";
+import { getResponseBuilder } from "./builders";
+import {
+  BaseSynthesizer,
+  ResponseBuilder,
+  SynthesizeParamsNonStreaming,
+  SynthesizeParamsStreaming,
+} from "./types";

 /**
 * A ResponseSynthesizer is used to generate a response from a query and a list of nodes.
 */
 export class ResponseSynthesizer implements BaseSynthesizer {
-  responseBuilder: BaseResponseBuilder;
+  responseBuilder: ResponseBuilder;
  serviceContext: ServiceContext;
  metadataMode: MetadataMode;

@@ -18,7 +23,7 @@ export class ResponseSynthesizer implements BaseSynthesizer {
    serviceContext,
    metadataMode = MetadataMode.NONE,
  }: {
-    responseBuilder?: BaseResponseBuilder;
+    responseBuilder?: ResponseBuilder;
    serviceContext?: ServiceContext;
    metadataMode?: MetadataMode;
  } = {}) {
@@ -28,22 +33,36 @@ export class ResponseSynthesizer implements BaseSynthesizer {
    this.metadataMode = metadataMode;
  }

-  async synthesize(
-    query: string,
-    nodesWithScore: NodeWithScore[],
-    parentEvent?: Event,
-  ) {
-    let textChunks: string[] = nodesWithScore.map(({ node }) =>
+  synthesize(
+    params: SynthesizeParamsStreaming,
+  ): Promise<AsyncIterable<Response>>;
+  synthesize(params: SynthesizeParamsNonStreaming): Promise<Response>;
+  async synthesize({
+    query,
+    nodesWithScore,
+    parentEvent,
+    stream,
+  }: SynthesizeParamsStreaming | SynthesizeParamsNonStreaming): Promise<
+    AsyncIterable<Response> | Response
+  > {
+    const textChunks: string[] = nodesWithScore.map(({ node }) =>
      node.getContent(this.metadataMode),
    );
-    const response = await this.responseBuilder.getResponse(
+    const nodes = nodesWithScore.map(({ node }) => node);
+    if (stream) {
+      const response = await this.responseBuilder.getResponse({
+        query,
+        textChunks,
+        parentEvent,
+        stream,
+      });
+      return streamConverter(response, (chunk) => new Response(chunk, nodes));
+    }
+    const response = await this.responseBuilder.getResponse({
      query,
      textChunks,
      parentEvent,
-    );
-    return new Response(
-      response,
-      nodesWithScore.map(({ node }) => node),
-    );
+    });
+    return new Response(response, nodes);
  }
 }
@@ -1,5 +1,6 @@
 import { Event } from "../callbacks/CallbackManager";
-import { LLM } from "../llm/LLM";
+import { LLM } from "../llm";
+import { streamConverter } from "../llm/utils";
 import {
  defaultRefinePrompt,
  defaultTextQaPrompt,
@@ -9,8 +10,13 @@ import {
  TextQaPrompt,
  TreeSummarizePrompt,
 } from "../Prompt";
-import { getBiggestPrompt } from "../PromptHelper";
+import { getBiggestPrompt, PromptHelper } from "../PromptHelper";
 import { ServiceContext } from "../ServiceContext";
+import {
+  ResponseBuilder,
+  ResponseBuilderParamsNonStreaming,
+  ResponseBuilderParamsStreaming,
+} from "./types";

 /**
 * Response modes of the response synthesizer
@@ -22,29 +28,10 @@ enum ResponseMode {
  SIMPLE = "simple",
 }

-/**
- * A ResponseBuilder is used in a response synthesizer to generate a response from multiple response chunks.
- */
-export interface BaseResponseBuilder {
-  /**
-   * Get the response from a query and a list of text chunks.
-   * @param query
-   * @param textChunks
-   * @param parentEvent
-   * @param prevResponse
-   */
-  getResponse(
-    query: string,
-    textChunks: string[],
-    parentEvent?: Event,
-    prevResponse?: string,
-  ): Promise<string>;
-}
-
 /**
 * A response builder that just concatenates responses.
 */
-export class SimpleResponseBuilder implements BaseResponseBuilder {
+export class SimpleResponseBuilder implements ResponseBuilder {
  llm: LLM;
  textQATemplate: SimplePrompt;

@@ -53,27 +40,42 @@ export class SimpleResponseBuilder implements BaseResponseBuilder {
    this.textQATemplate = defaultTextQaPrompt;
  }

-  async getResponse(
-    query: string,
-    textChunks: string[],
-    parentEvent?: Event,
-  ): Promise<string> {
+  getResponse(
+    params: ResponseBuilderParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  getResponse(params: ResponseBuilderParamsNonStreaming): Promise<string>;
+  async getResponse({
+    query,
+    textChunks,
+    parentEvent,
+    stream,
+  }:
+    | ResponseBuilderParamsStreaming
+    | ResponseBuilderParamsNonStreaming): Promise<
+    AsyncIterable<string> | string
+  > {
    const input = {
      query,
      context: textChunks.join("\n\n"),
    };

    const prompt = this.textQATemplate(input);
-    const response = await this.llm.complete(prompt, parentEvent);
-    return response.message.content;
+    if (stream) {
+      const response = await this.llm.complete({ prompt, parentEvent, stream });
+      return streamConverter(response, (chunk) => chunk.text);
+    } else {
+      const response = await this.llm.complete({ prompt, parentEvent, stream });
+      return response.text;
+    }
  }
 }

 /**
 * A response builder that uses the query to ask the LLM generate a better response using multiple text chunks.
 */
-export class Refine implements BaseResponseBuilder {
-  serviceContext: ServiceContext;
+export class Refine implements ResponseBuilder {
+  llm: LLM;
+  promptHelper: PromptHelper;
  textQATemplate: TextQaPrompt;
  refineTemplate: RefinePrompt;

@@ -82,31 +84,48 @@ export class Refine implements BaseResponseBuilder {
    textQATemplate?: TextQaPrompt,
    refineTemplate?: RefinePrompt,
  ) {
-    this.serviceContext = serviceContext;
+    this.llm = serviceContext.llm;
+    this.promptHelper = serviceContext.promptHelper;
    this.textQATemplate = textQATemplate ?? defaultTextQaPrompt;
    this.refineTemplate = refineTemplate ?? defaultRefinePrompt;
  }

-  async getResponse(
-    query: string,
-    textChunks: string[],
-    parentEvent?: Event,
-    prevResponse?: string,
-  ): Promise<string> {
-    let response: string | undefined = undefined;
+  getResponse(
+    params: ResponseBuilderParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  getResponse(params: ResponseBuilderParamsNonStreaming): Promise<string>;
+  async getResponse({
+    query,
+    textChunks,
+    parentEvent,
+    prevResponse,
+    stream,
+  }:
+    | ResponseBuilderParamsStreaming
+    | ResponseBuilderParamsNonStreaming): Promise<
+    AsyncIterable<string> | string
+  > {
+    let response: AsyncIterable<string> | string | undefined = prevResponse;

-    for (const chunk of textChunks) {
-      if (!prevResponse) {
-        response = await this.giveResponseSingle(query, chunk, parentEvent);
-      } else {
-        response = await this.refineResponseSingle(
-          prevResponse,
+    for (let i = 0; i < textChunks.length; i++) {
+      const chunk = textChunks[i];
+      const lastChunk = i === textChunks.length - 1;
+      if (!response) {
+        response = await this.giveResponseSingle(
          query,
          chunk,
+          !!stream && lastChunk,
+          parentEvent,
+        );
+      } else {
+        response = await this.refineResponseSingle(
+          response as string,
+          query,
+          chunk,
+          !!stream && lastChunk,
          parentEvent,
        );
      }
-      prevResponse = response;
    }

    return response ?? "Empty Response";
@@ -115,153 +134,204 @@ export class Refine implements BaseResponseBuilder {
  private async giveResponseSingle(
    queryStr: string,
    textChunk: string,
+    stream: boolean,
    parentEvent?: Event,
-  ): Promise<string> {
+  ) {
    const textQATemplate: SimplePrompt = (input) =>
      this.textQATemplate({ ...input, query: queryStr });
-    const textChunks = this.serviceContext.promptHelper.repack(textQATemplate, [
-      textChunk,
-    ]);
+    const textChunks = this.promptHelper.repack(textQATemplate, [textChunk]);

-    let response: string | undefined = undefined;
+    let response: AsyncIterable<string> | string | undefined = undefined;

-    for (const chunk of textChunks) {
+    for (let i = 0; i < textChunks.length; i++) {
+      const chunk = textChunks[i];
+      const lastChunk = i === textChunks.length - 1;
      if (!response) {
-        response = (
-          await this.serviceContext.llm.complete(
-            textQATemplate({
-              context: chunk,
-            }),
-            parentEvent,
-          )
-        ).message.content;
+        response = await this.complete({
+          prompt: textQATemplate({
+            context: chunk,
+          }),
+          parentEvent,
+          stream: stream && lastChunk,
+        });
      } else {
        response = await this.refineResponseSingle(
-          response,
+          response as string,
          queryStr,
          chunk,
+          stream && lastChunk,
          parentEvent,
        );
      }
    }

-    return response ?? "Empty Response";
+    return response;
  }

  private async refineResponseSingle(
-    response: string,
+    initialReponse: string,
    queryStr: string,
    textChunk: string,
+    stream: boolean,
    parentEvent?: Event,
  ) {
    const refineTemplate: SimplePrompt = (input) =>
      this.refineTemplate({ ...input, query: queryStr });

-    const textChunks = this.serviceContext.promptHelper.repack(refineTemplate, [
-      textChunk,
-    ]);
+    const textChunks = this.promptHelper.repack(refineTemplate, [textChunk]);

-    for (const chunk of textChunks) {
-      response = (
-        await this.serviceContext.llm.complete(
-          refineTemplate({
-            context: chunk,
-            existingAnswer: response,
-          }),
-          parentEvent,
-        )
-      ).message.content;
+    let response: AsyncIterable<string> | string = initialReponse;
+
+    for (let i = 0; i < textChunks.length; i++) {
+      const chunk = textChunks[i];
+      const lastChunk = i === textChunks.length - 1;
+      response = await this.complete({
+        prompt: refineTemplate({
+          context: chunk,
+          existingAnswer: response as string,
+        }),
+        parentEvent,
+        stream: stream && lastChunk,
+      });
    }
    return response;
  }
+
+  async complete(params: {
+    prompt: string;
+    stream: boolean;
+    parentEvent?: Event;
+  }): Promise<AsyncIterable<string> | string> {
+    if (params.stream) {
+      const response = await this.llm.complete({ ...params, stream: true });
+      return streamConverter(response, (chunk) => chunk.text);
+    } else {
+      const response = await this.llm.complete({ ...params, stream: false });
+      return response.text;
+    }
+  }
 }

 /**
 * CompactAndRefine is a slight variation of Refine that first compacts the text chunks into the smallest possible number of chunks.
 */
 export class CompactAndRefine extends Refine {
-  async getResponse(
-    query: string,
-    textChunks: string[],
-    parentEvent?: Event,
-    prevResponse?: string,
-  ): Promise<string> {
+  getResponse(
+    params: ResponseBuilderParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  getResponse(params: ResponseBuilderParamsNonStreaming): Promise<string>;
+  async getResponse({
+    query,
+    textChunks,
+    parentEvent,
+    prevResponse,
+    stream,
+  }:
+    | ResponseBuilderParamsStreaming
+    | ResponseBuilderParamsNonStreaming): Promise<
+    AsyncIterable<string> | string
+  > {
    const textQATemplate: SimplePrompt = (input) =>
      this.textQATemplate({ ...input, query: query });
    const refineTemplate: SimplePrompt = (input) =>
      this.refineTemplate({ ...input, query: query });

    const maxPrompt = getBiggestPrompt([textQATemplate, refineTemplate]);
-    const newTexts = this.serviceContext.promptHelper.repack(
-      maxPrompt,
-      textChunks,
-    );
-    const response = super.getResponse(
+    const newTexts = this.promptHelper.repack(maxPrompt, textChunks);
+    const params = {
      query,
-      newTexts,
+      textChunks: newTexts,
      parentEvent,
      prevResponse,
-    );
-    return response;
+    };
+    if (stream) {
+      return super.getResponse({
+        ...params,
+        stream,
+      });
+    }
+    return super.getResponse(params);
  }
 }
+
 /**
 * TreeSummarize repacks the text chunks into the smallest possible number of chunks and then summarizes them, then recursively does so until there's one chunk left.
 */
-export class TreeSummarize implements BaseResponseBuilder {
-  serviceContext: ServiceContext;
+export class TreeSummarize implements ResponseBuilder {
+  llm: LLM;
+  promptHelper: PromptHelper;
  summaryTemplate: TreeSummarizePrompt;

  constructor(
    serviceContext: ServiceContext,
    summaryTemplate?: TreeSummarizePrompt,
  ) {
-    this.serviceContext = serviceContext;
+    this.llm = serviceContext.llm;
+    this.promptHelper = serviceContext.promptHelper;
    this.summaryTemplate = summaryTemplate ?? defaultTreeSummarizePrompt;
  }

-  async getResponse(
-    query: string,
-    textChunks: string[],
-    parentEvent?: Event,
-  ): Promise<string> {
+  getResponse(
+    params: ResponseBuilderParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  getResponse(params: ResponseBuilderParamsNonStreaming): Promise<string>;
+  async getResponse({
+    query,
+    textChunks,
+    parentEvent,
+    stream,
+  }:
+    | ResponseBuilderParamsStreaming
+    | ResponseBuilderParamsNonStreaming): Promise<
+    AsyncIterable<string> | string
+  > {
    if (!textChunks || textChunks.length === 0) {
      throw new Error("Must have at least one text chunk");
    }

    // Should we send the query here too?
-    const packedTextChunks = this.serviceContext.promptHelper.repack(
+    const packedTextChunks = this.promptHelper.repack(
      this.summaryTemplate,
      textChunks,
    );

    if (packedTextChunks.length === 1) {
-      return (
-        await this.serviceContext.llm.complete(
-          this.summaryTemplate({
-            context: packedTextChunks[0],
-            query,
-          }),
-          parentEvent,
-        )
-      ).message.content;
+      const params = {
+        prompt: this.summaryTemplate({
+          context: packedTextChunks[0],
+          query,
+        }),
+        parentEvent,
+      };
+      if (stream) {
+        const response = await this.llm.complete({ ...params, stream });
+        return streamConverter(response, (chunk) => chunk.text);
+      }
+      return (await this.llm.complete(params)).text;
    } else {
      const summaries = await Promise.all(
        packedTextChunks.map((chunk) =>
-          this.serviceContext.llm.complete(
-            this.summaryTemplate({
+          this.llm.complete({
+            prompt: this.summaryTemplate({
              context: chunk,
              query,
            }),
            parentEvent,
-          ),
+          }),
        ),
      );

-      return this.getResponse(
+      const params = {
        query,
-        summaries.map((s) => s.message.content),
-      );
+        textChunks: summaries.map((s) => s.text),
+      };
+      if (stream) {
+        return this.getResponse({
+          ...params,
+          stream,
+        });
+      }
+      return this.getResponse(params);
    }
  }
 }
@@ -269,7 +339,7 @@ export class TreeSummarize implements BaseResponseBuilder {
 export function getResponseBuilder(
  serviceContext: ServiceContext,
  responseMode?: ResponseMode,
-): BaseResponseBuilder {
+): ResponseBuilder {
  switch (responseMode) {
    case ResponseMode.SIMPLE:
      return new SimpleResponseBuilder(serviceContext);
@@ -2,14 +2,57 @@ import { Event } from "../callbacks/CallbackManager";
 import { NodeWithScore } from "../Node";
 import { Response } from "../Response";

+export interface SynthesizeParamsBase {
+  query: string;
+  nodesWithScore: NodeWithScore[];
+  parentEvent?: Event;
+}
+
+export interface SynthesizeParamsStreaming extends SynthesizeParamsBase {
+  stream: true;
+}
+
+export interface SynthesizeParamsNonStreaming extends SynthesizeParamsBase {
+  stream?: false | null;
+}
+
 /**
 * A BaseSynthesizer is used to generate a response from a query and a list of nodes.
- * TODO: convert response builders to implement this interface (similar to Python).
 */
 export interface BaseSynthesizer {
  synthesize(
-    query: string,
-    nodesWithScore: NodeWithScore[],
-    parentEvent?: Event,
-  ): Promise<Response>;
+    params: SynthesizeParamsStreaming,
+  ): Promise<AsyncIterable<Response>>;
+  synthesize(params: SynthesizeParamsNonStreaming): Promise<Response>;
+}
+
+export interface ResponseBuilderParamsBase {
+  query: string;
+  textChunks: string[];
+  parentEvent?: Event;
+  prevResponse?: string;
+}
+
+export interface ResponseBuilderParamsStreaming
+  extends ResponseBuilderParamsBase {
+  stream: true;
+}
+
+export interface ResponseBuilderParamsNonStreaming
+  extends ResponseBuilderParamsBase {
+  stream?: false | null;
+}
+
+/**
+ * A ResponseBuilder is used in a response synthesizer to generate a response from multiple response chunks.
+ */
+export interface ResponseBuilder {
+  /**
+   * Get the response from a query and a list of text chunks.
+   * @param params
+   */
+  getResponse(
+    params: ResponseBuilderParamsStreaming,
+  ): Promise<AsyncIterable<string>>;
+  getResponse(params: ResponseBuilderParamsNonStreaming): Promise<string>;
 }
@@ -1,7 +1,7 @@
-import { CallbackManager, Event } from "../../callbacks/CallbackManager";
+import { CallbackManager } from "../../callbacks/CallbackManager";
 import { OpenAIEmbedding } from "../../embeddings";
 import { globalsHelper } from "../../GlobalsHelper";
-import { ChatMessage, OpenAI } from "../../llm/LLM";
+import { LLMChatParamsBase, OpenAI } from "../../llm/LLM";

 export function mockLlmGeneration({
  languageModel,
@@ -13,7 +13,7 @@ export function mockLlmGeneration({
  jest
    .spyOn(languageModel, "chat")
    .mockImplementation(
-      async (messages: ChatMessage[], parentEvent?: Event) => {
+      async ({ messages, parentEvent }: LLMChatParamsBase) => {
        const text = "MOCK_TOKEN_1-MOCK_TOKEN_2";
        const event = globalsHelper.createEvent({
          parentEvent,
@@ -162,6 +162,10 @@ export const installTemplate = async (
      props.openAiKey,
      props.vectorDb,
    );
+  } else {
+    // this is a frontend for a full-stack app, create .env file with model information
+    const content = `MODEL=${props.model}\nNEXT_PUBLIC_MODEL=${props.model}\n`;
+    await fs.writeFile(path.join(props.root, ".env"), content);
  }
 };

@@ -1,17 +1,43 @@
 import {
+  JSONValue,
  createCallbacksTransformer,
  createStreamDataTransformer,
+  experimental_StreamData,
  trimStartOfStreamHelper,
  type AIStreamCallbacksAndOptions,
 } from "ai";

-function createParser(res: AsyncGenerator<any>) {
+type ParserOptions = {
+  image_url?: string;
+};
+
+function createParser(
+  res: AsyncGenerator<any>,
+  data: experimental_StreamData,
+  opts?: ParserOptions,
+) {
  const trimStartOfStream = trimStartOfStreamHelper();
  return new ReadableStream<string>({
+    start() {
+      // if image_url is provided, send it via the data stream
+      if (opts?.image_url) {
+        const message: JSONValue = {
+          type: "image_url",
+          image_url: {
+            url: opts.image_url,
+          },
+        };
+        data.append(message);
+      } else {
+        data.append({}); // send an empty image response for the user's message
+      }
+    },
    async pull(controller): Promise<void> {
      const { value, done } = await res.next();
      if (done) {
        controller.close();
+        data.append({}); // send an empty image response for the assistant's message
+        data.close();
        return;
      }

@@ -25,11 +51,16 @@ function createParser(res: AsyncGenerator<any>) {

 export function LlamaIndexStream(
  res: AsyncGenerator<any>,
-  callbacks?: AIStreamCallbacksAndOptions,
-): ReadableStream {
-  return createParser(res)
-    .pipeThrough(createCallbacksTransformer(callbacks))
-    .pipeThrough(
-      createStreamDataTransformer(callbacks?.experimental_streamData),
-    );
+  opts?: {
+    callbacks?: AIStreamCallbacksAndOptions;
+    parserOptions?: ParserOptions;
+  },
+): { stream: ReadableStream; data: experimental_StreamData } {
+  const data = new experimental_StreamData();
+  return {
+    stream: createParser(res, data, opts?.parserOptions)
+      .pipeThrough(createCallbacksTransformer(opts?.callbacks))
+      .pipeThrough(createStreamDataTransformer(true)),
+    data,
+  };
 }
@@ -1,5 +1,5 @@
 import { Message, StreamingTextResponse } from "ai";
-import { MessageContent, OpenAI } from "llamaindex";
+import { ChatMessage, MessageContent, OpenAI } from "llamaindex";
 import { NextRequest, NextResponse } from "next/server";
 import { createChatEngine } from "./engine";
 import { LlamaIndexStream } from "./llamaindex-stream";
@@ -42,7 +42,7 @@ export async function POST(request: NextRequest) {
    }

    const llm = new OpenAI({
-      model: process.env.MODEL || "gpt-3.5-turbo",
+      model: (process.env.MODEL as any) ?? "gpt-3.5-turbo",
      maxTokens: 512,
    });

@@ -55,15 +55,19 @@ export async function POST(request: NextRequest) {

    const response = await chatEngine.chat(
      lastMessageContent as MessageContent,
-      messages,
+      messages as ChatMessage[],
      true,
    );

    // Transform the response into a readable stream
-    const stream = LlamaIndexStream(response);
+    const { stream, data: streamData } = LlamaIndexStream(response, {
+      parserOptions: {
+        image_url: data?.imageUrl,
+      },
+    });

    // Return a StreamingTextResponse, which can be consumed by the client
-    return new StreamingTextResponse(stream);
+    return new StreamingTextResponse(stream, {}, streamData);
  } catch (error) {
    console.error("[LlamaIndex]", error);
    return NextResponse.json(
@@ -1,6 +1,8 @@
 "use client";

 import { useChat } from "ai/react";
+import { useMemo } from "react";
+import { insertDataIntoMessages } from "./transform";
 import { ChatInput, ChatMessages } from "./ui/chat";

 export default function ChatSection() {
@@ -12,6 +14,7 @@ export default function ChatSection() {
    handleInputChange,
    reload,
    stop,
+    data,
  } = useChat({
    api: process.env.NEXT_PUBLIC_CHAT_API,
    headers: {
@@ -19,10 +22,14 @@ export default function ChatSection() {
    },
  });

+  const transformedMessages = useMemo(() => {
+    return insertDataIntoMessages(messages, data);
+  }, [messages, data]);
+
  return (
    <div className="space-y-4 max-w-5xl w-full">
      <ChatMessages
-        messages={messages}
+        messages={transformedMessages}
        isLoading={isLoading}
        reload={reload}
        stop={stop}
@@ -0,0 +1,19 @@
+import { JSONValue, Message } from "ai";
+
+export const isValidMessageData = (rawData: JSONValue | undefined) => {
+  if (!rawData || typeof rawData !== "object") return false;
+  if (Object.keys(rawData).length === 0) return false;
+  return true;
+};
+
+export const insertDataIntoMessages = (
+  messages: Message[],
+  data: JSONValue[] | undefined,
+) => {
+  if (!data) return messages;
+  messages.forEach((message, i) => {
+    const rawData = data[i];
+    if (isValidMessageData(rawData)) message.data = rawData;
+  });
+  return messages;
+};
@@ -1,18 +1,49 @@
 import { Check, Copy } from "lucide-react";

+import { JSONValue, Message } from "ai";
+import Image from "next/image";
 import { Button } from "../button";
 import ChatAvatar from "./chat-avatar";
-import { Message } from "./chat.interface";
 import Markdown from "./markdown";
 import { useCopyToClipboard } from "./use-copy-to-clipboard";

+interface ChatMessageImageData {
+  type: "image_url";
+  image_url: {
+    url: string;
+  };
+}
+
+// This component will parse message data and render the appropriate UI.
+function ChatMessageData({ messageData }: { messageData: JSONValue }) {
+  const { image_url, type } = messageData as unknown as ChatMessageImageData;
+  if (type === "image_url") {
+    return (
+      <div className="rounded-md max-w-[200px] shadow-md">
+        <Image
+          src={image_url.url}
+          width={0}
+          height={0}
+          sizes="100vw"
+          style={{ width: "100%", height: "auto" }}
+          alt=""
+        />
+      </div>
+    );
+  }
+  return null;
+}
+
 export default function ChatMessage(chatMessage: Message) {
  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
  return (
    <div className="flex items-start gap-4 pr-5 pt-5">
      <ChatAvatar role={chatMessage.role} />
      <div className="group flex flex-1 justify-between gap-2">
-        <div className="flex-1">
+        <div className="flex-1 space-y-4">
+          {chatMessage.data && (
+            <ChatMessageData messageData={chatMessage.data} />
+          )}
          <Markdown content={chatMessage.content} />
        </div>
        <Button
@@ -1,8 +1,4 @@
-export interface Message {
-  id: string;
-  content: string;
-  role: string;
-}
+import { Message } from "ai";

 export interface ChatHandler {
  messages: Message[];
@@ -1,5 +1,5 @@
 import ChatInput from "./chat-input";
 import ChatMessages from "./chat-messages";

-export { type ChatHandler, type Message } from "./chat.interface";
+export { type ChatHandler } from "./chat.interface";
 export { ChatInput, ChatMessages };
Author	SHA1	Message	Date
Alex Yang	0f47d185c3	feat: code	2024-01-15 12:52:25 -06:00
Alex Yang	8e2beaddca	Merge remote-tracking branch 'origin/main' into himself65/fix-type # Conflicts: # packages/core/src/llm/LLM.ts	2024-01-15 12:49:51 -06:00
Alex Yang	bd3a7fd450	feat: improve	2024-01-15 12:48:32 -06:00
Marcus Schiesser	2001eb7ffb	fix: format	2024-01-15 18:15:44 +07:00
Marcus Schiesser	f18c9f69d4	refactor: update low-level streaming interface (#325 )	2024-01-15 18:06:53 +07:00
Thuc Pham	8e124e5b63	feat: support showing image for chat message in NextJS (#368 )	2024-01-15 17:57:20 +07:00
Alex Yang	2f1894b251	docs(changeset): fix: openai type might be missing	2024-01-14 21:32:18 -06:00
Alex Yang	2f1afecea7	fix: abstract openai class	2024-01-14 21:30:54 -06:00
Nir Gazit	4ed5e544b0	docs: added openllmetry observability (#369 )	2024-01-15 10:15:02 +07:00