fix: fix and test

opendeeple · Dec 17, 2024 · bd9a22e · bd9a22e
1 parent e60cc9d
commit bd9a22e
Show file tree

Hide file tree

Showing 6 changed files with 86 additions and 65 deletions.
diff --git a/src/.DS_Store b/src/.DS_Store
diff --git a/src/anthropic/index.ts b/src/anthropic/index.ts
@@ -1,14 +1,11 @@
 import Anthropic from "@anthropic-ai/sdk";
-import { omit } from "../utils";
-import {
-  AnthropicChatCompletationParams,
-  FetchAIChatCompletation,
-} from "../type";
+import { omit } from "../utils/index";
+import { AnthropicChatCompletionParams, FetchAIChatCompletion } from "../type";
 
 export default class AnthropicChatRepository {
   constructor(readonly provider: Anthropic) {}
 
-  async create(body: AnthropicChatCompletationParams) {
+  async create(body: AnthropicChatCompletionParams) {
     try {
       const completion = await this.provider.messages.create(
         omit(body, ["prediction_tokens"])
@@ -18,7 +15,7 @@ export default class AnthropicChatRepository {
         content.type == "text"
           ? content.text
           : `${content.name}(id=${content.id}, value=${content.input})`;
-      const result: FetchAIChatCompletation = {
+      const result: FetchAIChatCompletion = {
         provider: "Anthropic",
         success: true,
         prediction: {
@@ -39,7 +36,7 @@ export default class AnthropicChatRepository {
       return result;
     } catch (error) {
       const message = error instanceof Error ? error.message : String(error);
-      const result: FetchAIChatCompletation = {
+      const result: FetchAIChatCompletion = {
         provider: "Anthropic",
         success: false,
         error: message,
@@ -54,7 +51,7 @@ export default class AnthropicChatRepository {
     }
   }
 
-  async countTokens(body: AnthropicChatCompletationParams) {
+  async countTokens(body: AnthropicChatCompletionParams) {
     const response = await this.provider.beta.messages.countTokens({
       betas: ["token-counting-2024-11-01"],
       ...omit(body, ["max_tokens"]),

diff --git a/src/fetchai/index.ts b/src/fetchai/index.ts
@@ -3,9 +3,12 @@ import Anthropic from "@anthropic-ai/sdk";
 import AnthropicChatRepository from "../anthropic";
 import OpenAIChatRepository from "../openai";
 import {
-  AnthropicChatCompletationParams,
+  AnthropicChatCompletionParams,
   ClientOptions,
-  FetchAIChatCompletationParams,
+  FetchAIChatCompletionParams,
+  FetchAIChatCount,
+  FetchAIChatModels,
+  FetchAIProviders,
 } from "../type";
 
 export default class FetchAIChatRepository {
@@ -19,23 +22,39 @@ export default class FetchAIChatRepository {
     );
   }
 
+  private isAnthropicModel(model: FetchAIChatModels) {
+    return /^claude/.test(model);
+  }
+
   private isAnthropic(
-    body: FetchAIChatCompletationParams
-  ): body is AnthropicChatCompletationParams {
-    return /^claude/.test(body.model);
+    body: FetchAIChatCompletionParams
+  ): body is AnthropicChatCompletionParams {
+    return this.isAnthropicModel(body.model);
   }
 
-  async create(body: FetchAIChatCompletationParams) {
+  async create(body: FetchAIChatCompletionParams) {
     if (this.isAnthropic(body)) {
       return this.Anthropic.create(body);
     }
     return this.OpenAI.create(body);
   }
 
-  async countTokens(body: FetchAIChatCompletationParams) {
+  async countTokens(
+    body: FetchAIChatCompletionParams
+  ): Promise<FetchAIChatCount> {
     if (this.isAnthropic(body)) {
-      return this.Anthropic.countTokens(body);
+      return {
+        provider: "Anthropic",
+        input_tokens: await this.Anthropic.countTokens(body),
+      };
     }
-    return this.OpenAI.countTokens(body);
+    return {
+      provider: "OpenAI",
+      input_tokens: await this.OpenAI.countTokens(body),
+    };
+  }
+
+  provider(model: FetchAIChatModels): FetchAIProviders {
+    return this.isAnthropicModel(model) ? "Anthropic" : "OpenAI";
   }
 }
diff --git a/src/index.ts b/src/index.ts
@@ -1,5 +1,5 @@
-import AnthropicChatRepository from "./anthropic";
-import FetchAIChatRepository from "./fetchai";
-import OpenAIChatRepository from "./openai";
+import AnthropicChatRepository from "./anthropic/index";
+import FetchAIChatRepository from "./fetchai/index";
+import OpenAIChatRepository from "./openai/index";
 export { OpenAIChatRepository, AnthropicChatRepository };
 export default FetchAIChatRepository;
diff --git a/src/openai/index.ts b/src/openai/index.ts
@@ -1,15 +1,11 @@
 import OpenAI from "openai";
 import { encoding_for_model } from "tiktoken";
-import {
-  FetchAIChatCompletation,
-  OpenAIChatCompletationParams,
-  OpenAIChatMessageParams,
-} from "../type";
+import { FetchAIChatCompletion, OpenAIChatCompletionParams } from "../type";
 
 export default class OpenAIChatRepository {
   constructor(readonly provider: OpenAI) {}
 
-  async create(body: OpenAIChatCompletationParams) {
+  async create(body: OpenAIChatCompletionParams) {
     const prediction = body.prediction;
     try {
       const completion = await this.provider.chat.completions.create({
@@ -23,7 +19,7 @@ export default class OpenAIChatRepository {
             }
           : undefined,
       });
-      const result: FetchAIChatCompletation = {
+      const result: FetchAIChatCompletion = {
         provider: "OpenAI",
         success: true,
         prediction: {
@@ -43,7 +39,7 @@ export default class OpenAIChatRepository {
       return result;
     } catch (error) {
       const message = error instanceof Error ? error.message : String(error);
-      const result: FetchAIChatCompletation = {
+      const result: FetchAIChatCompletion = {
         provider: "OpenAI",
         success: false,
         error: message,
@@ -58,10 +54,10 @@ export default class OpenAIChatRepository {
     }
   }
 
-  async countTokens(body: OpenAIChatCompletationParams) {
+  async countTokens(body: OpenAIChatCompletionParams) {
     const modelEncoding = encoding_for_model(body.model);
 
-    let tokens = modelEncoding.encode(body.system).length;
+    let tokens = body.system ? modelEncoding.encode(body.system).length : 0;
 
     const messages = body.messages;
     for (const message of messages) {
@@ -80,14 +76,14 @@ export default class OpenAIChatRepository {
     return tokens;
   }
 
-  private buildMessages(body: OpenAIChatCompletationParams) {
-    const messages: OpenAIChatMessageParams[] = [
-      {
+  private buildMessages(body: OpenAIChatCompletionParams) {
+    const messages: OpenAI.Chat.ChatCompletionMessageParam[] = [];
+    if (body.system) {
+      messages.push({
         role: "system",
         content: body.system,
-      },
-      ...body.messages,
-    ];
-    return messages;
+      });
+    }
+    return [...messages, ...body.messages];
   }
 }
diff --git a/src/type.ts b/src/type.ts
@@ -6,23 +6,29 @@ import { TiktokenModel } from "tiktoken";
 
 export type Maybe<T> = T | null | undefined;
 
+export interface FetchAIChatMessageParams {
+  role: "assistant" | "user";
+  content: string;
+}
+
 export type AnthropicChatModels = Anthropic.Messages.Model;
-export type AnthropicChatMessageParams = Anthropic.Messages.MessageParam;
+export type AnthropicChatMessageParams = FetchAIChatMessageParams;
 
 export type AnthropicChatCompletionParams = {
   model: AnthropicChatModels;
-  system?: string; // Optional system message
+  system?: string;
   messages: AnthropicChatMessageParams[];
+  prediction?: string;
   max_tokens: number;
   prediction_tokens?: number;
 };
 
 export type OpenAIChatModels = TiktokenModel;
-export type OpenAIChatMessageParams = OpenAI.Chat.ChatCompletionMessageParam;
+export type OpenAIChatMessageParams = FetchAIChatMessageParams;
 
 export type OpenAIChatCompletionParams = {
   model: OpenAIChatModels;
-  system?: string; // Optional system message
+  system?: string;
   messages: OpenAIChatMessageParams[];
   prediction?: string;
   max_tokens: number;
@@ -34,36 +40,39 @@ export interface ClientOptions {
   anthropic?: AnthropicClientOptions;
 }
 
-export type FetchAIChatMessageParams =
-  | (OpenAIChatMessageParams & {
-      role: "system" | "user" | "assistant";
-    })
-  | (AnthropicChatMessageParams & {
-      role: "user" | "assistant";
-    });
-
 export type FetchAIChatModels = OpenAIChatModels | AnthropicChatModels;
 
 export type FetchAIChatCompletionParams =
   | OpenAIChatCompletionParams
   | AnthropicChatCompletionParams;
 
+export type FetchAIProviders = "OpenAI" | "Anthropic";
+
+export interface FetchAIChatPrediction {
+  id?: string;
+  content: string;
+  openai?: OpenAI.Chat.Completions.ChatCompletion;
+  anthropic?: Omit<Anthropic.Messages.Message, "id">;
+}
+
+export interface FetchAIUsage {
+  input_tokens?: number;
+  output_tokens?: number;
+  total_tokens?: number;
+  predicted_tokens: number;
+  openai?: OpenAI.Completions.CompletionUsage;
+  anthropic?: Anthropic.Messages.Usage;
+}
+
 export interface FetchAIChatCompletion {
-  provider: "OpenAI" | "Anthropic";
+  provider: FetchAIProviders;
   success: boolean;
   error?: string;
-  prediction: {
-    id?: string;
-    content: string;
-    openai?: OpenAI.Chat.Completions.ChatCompletion;
-    anthropic?: Omit<Anthropic.Messages.Message, "id">;
-  };
-  usage: {
-    input_tokens?: number;
-    output_tokens?: number;
-    total_tokens?: number;
-    predicted_tokens: number;
-    openai?: OpenAI.Completions.CompletionUsage;
-    anthropic?: Anthropic.Messages.Usage;
-  };
+  prediction: FetchAIChatPrediction;
+  usage: FetchAIUsage;
+}
+
+export interface FetchAIChatCount {
+  provider: FetchAIProviders;
+  input_tokens: number;
 }