Add custom model support (#15)

mme · Nov 22, 2023 · fb0f7fa · fb0f7fa
1 parent 28b31e8
commit fb0f7fa
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 10 deletions.
diff --git a/packages/core/src/beak.ts b/packages/core/src/beak.ts
@@ -1,5 +1,5 @@
 import { EventEmitter } from "eventemitter3";
-import { OpenAI, OpenAIModel } from "@beakjs/openai";
+import { OpenAI, OpenAIModel, CustomModel } from "@beakjs/openai";
 import {
   LLMAdapter,
   Message,
@@ -22,7 +22,7 @@ const FORMATTING_INSTRUCTIONS =
 export interface BeakConfiguration {
   openAIApiKey?: string;
   baseUrl?: string;
-  openAIModel?: OpenAIModel;
+  openAIModel?: OpenAIModel | CustomModel;
   maxFeedback?: number;
   instructions?: string;
   temperature?: number;

diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts
@@ -1,4 +1,4 @@
 export { BeakCore } from "./beak";
 export type { FunctionDefinition } from "./types";
 export { Message, DebugLogger } from "./types";
-export type { OpenAIModel } from "@beakjs/openai";
+export type { OpenAIModel, CustomModel } from "@beakjs/openai";
diff --git a/packages/openai/src/index.ts b/packages/openai/src/index.ts
@@ -1,4 +1,9 @@
 export { OpenAI } from "./openai";
 export { ChatCompletion } from "./chat";
-export type { OpenAIModel, OpenAIMessage, OpenAIFunction } from "./types";
+export type {
+  OpenAIModel,
+  OpenAIMessage,
+  OpenAIFunction,
+  CustomModel,
+} from "./types";
 export type { FetchChatCompletionParams } from "./chat";
diff --git a/packages/openai/src/openai.ts b/packages/openai/src/openai.ts
@@ -8,13 +8,14 @@ import {
   DebugLogger,
   NoopDebugLogger,
   DEFAULT_MODEL,
+  CustomModel,
 } from "./types";
 import { ChatCompletion, FetchChatCompletionParams } from "./chat";
 
 interface OpenAIConfiguration {
   apiKey?: string;
   baseUrl?: string;
-  model?: OpenAIModel;
+  model?: OpenAIModel | CustomModel;
   debugLogger?: DebugLogger;
 }
 
@@ -32,7 +33,7 @@ interface OpenAIEvents {
 export class OpenAI extends EventEmitter<OpenAIEvents> {
   private apiKey?: string;
   private baseUrl?: string;
-  private model: OpenAIModel;
+  private model: OpenAIModel | CustomModel;
   private debug: DebugLogger;
 
   private completionClient: ChatCompletion | null = null;
@@ -50,9 +51,14 @@ export class OpenAI extends EventEmitter<OpenAIEvents> {
 
   public async queryChatCompletion(params: FetchChatCompletionParams) {
     params = { ...params };
-    params.maxTokens ||= maxTokensForModel(this.model);
+    if (!(this.model instanceof CustomModel)) {
+      params.maxTokens ||= maxTokensForModel(this.model);
+    } else if (!params.maxTokens) {
+      throw new Error("maxTokens must be specified for custom models.");
+    }
     params.functions ||= [];
-    params.model = this.model;
+    params.model =
+      this.model instanceof CustomModel ? this.model.name : this.model;
     params.messages = this.buildPrompt(params);
     return await this.runPrompt(params);
   }

diff --git a/packages/openai/src/types.ts b/packages/openai/src/types.ts
@@ -13,6 +13,10 @@ export type OpenAIModel =
   | "gpt-4-32k-0613"
   | "gpt-3.5-turbo-16k-0613";
 
+export class CustomModel {
+  constructor(public name: string) {}
+}
+
 export interface OpenAIChatCompletionChunk {
   choices: {
     delta: {

diff --git a/packages/react/src/Beak.tsx b/packages/react/src/Beak.tsx
@@ -1,5 +1,5 @@
 import React, { useMemo } from "react";
-import { BeakCore, OpenAIModel, DebugLogger } from "@beakjs/core";
+import { BeakCore, OpenAIModel, CustomModel, DebugLogger } from "@beakjs/core";
 import * as DefaultIcons from "./Icons";
 
 const DEFAULT_DEBUG_LOGGER = new DebugLogger([]);
@@ -49,7 +49,7 @@ export function useBeakContext(): BeakContext {
 interface BeakProps {
   __unsafeOpenAIApiKey__?: string;
   baseUrl?: string;
-  openAIModel?: OpenAIModel;
+  openAIModel?: OpenAIModel | CustomModel;
   temperature?: number;
   instructions?: string;
   maxFeedback?: number;