Pro 140/tokens limits (#77)

* Create TokenUsage model

* Log token after AI request

* Install Tremor and add placeholder bar chart

* wip

* create a models directory inside db

* Just use progress chart

* Add token summaries for Pro and Basic tokens

* wip

* Add model function for price

* wip - creating Custom pricing

* Add more states to Billing

* Make it possible to subscript to custom tier through dash

* Finish configuring billing page

* Throw error is user exceeds their token limit

* Add model flag to CLI

* Fix some tsc issues

* Improve billing states

* Cleanup

* Add changeset

* Bump package versions

* Name change to team plan

* Set rate limiting for Team tier

.changeset/popular-yaks-sing.md

@@ -0,0 +1,13 @@
+---
+"@floe/review-action": minor
+"@floe/features": minor
+"@floe/requests": minor
+"@floe/cli": minor
+"@floe/lib": minor
+"@floe/db": minor
+"@floe/ui": minor
+"@floe/api": minor
+"@floe/app": minor
+---
+
+Add support for token usage and pro / basic models.

.changeset/pre.json

@@ -29,6 +29,7 @@
     "hot-kangaroos-dream",
     "odd-phones-shave",
     "popular-rules-beam",
+    "popular-yaks-sing",
     "quick-rice-care",
     "silver-pumas-dance",
     "slimy-bikes-divide",

actions/review-action/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @floe/review-action
 
+## 0.1.0-beta.7
+
+### Minor Changes
+
+- 5f84851: Add support for token usage and pro / basic models.
+
+### Patch Changes
+
+- Updated dependencies [5f84851]
+  - @floe/features@0.1.0-beta.5
+  - @floe/requests@0.1.0-beta.6
+  - @floe/lib@0.1.0-beta.4
+
 ## 0.1.0-beta.6
 
 ### Minor Changes

actions/review-action/dist/index.js

@@ -45334,19 +45334,21 @@ const api = lib_axios.create({
 const querySchema = z.object({
     path: z.string(),
     content: z.string(),
-    startLine: z.coerce.number().optional().default(1),
+    startLine: z.coerce.number().default(1),
     rule: z.object({
         code: z.string(),
         level: z.union([z.literal("error"), z.literal("warn")]),
         description: z.string(),
     }),
+    model: z.union([z.literal("pro"), z.literal("basic")]).default("pro"),
 });
-async function createReview({ path, content, startLine, rule, }) {
+async function createReview({ path, content, startLine, rule, model, }) {
     return api.post("/api/v1/review", {
         path,
         content,
         startLine,
         rule,
+        model,
     });
 }
 
@@ -45390,14 +45392,15 @@ function checkIfUnderEvaluationLimit(evalutationsByFile, limit) {
  * Generate a review for each hunk and rule.
  * Output is an array of reviews grouped by file.
  */
-async function getReviewsByFile(evalutationsByFile) {
+async function getReviewsByFile(evalutationsByFile, options) {
     return Promise.all(evalutationsByFile.map(async ({ path, evaluations }) => {
         const evaluationsResponse = await Promise.all(evaluations.map(async ({ rule, hunk }) => {
             const review = await createReview({
                 path,
                 content: hunk.content,
                 startLine: hunk.startLine,
                 rule,
+                model: options?.model,
             });
             return {
                 review: {

actions/review-action/dist/packages/requests/review/_post.d.ts

@@ -3,7 +3,7 @@ import type OpenAI from "openai";
 export declare const querySchema: z.ZodObject<{
     path: z.ZodString;
     content: z.ZodString;
-    startLine: z.ZodDefault<z.ZodOptional<z.ZodNumber>>;
+    startLine: z.ZodDefault<z.ZodNumber>;
     rule: z.ZodObject<{
         code: z.ZodString;
         level: z.ZodUnion<[z.ZodLiteral<"error">, z.ZodLiteral<"warn">]>;
@@ -17,6 +17,7 @@ export declare const querySchema: z.ZodObject<{
         level: "error" | "warn";
         description: string;
     }>;
+    model: z.ZodDefault<z.ZodUnion<[z.ZodLiteral<"pro">, z.ZodLiteral<"basic">]>>;
 }, "strip", z.ZodTypeAny, {
     path: string;
     content: string;
@@ -26,6 +27,7 @@ export declare const querySchema: z.ZodObject<{
         level: "error" | "warn";
         description: string;
     };
+    model: "pro" | "basic";
 }, {
     path: string;
     content: string;
@@ -35,6 +37,7 @@ export declare const querySchema: z.ZodObject<{
         description: string;
     };
     startLine?: number | undefined;
+    model?: "pro" | "basic" | undefined;
 }>;
 export type PostReviewResponse = {
     violations: {
@@ -56,6 +59,6 @@ export type PostReviewResponse = {
     model: string;
     usage: OpenAI.Completions.CompletionUsage | undefined;
 } | undefined;
-export type PostReviewInput = z.infer<typeof querySchema>;
-export declare function createReview({ path, content, startLine, rule, }: PostReviewInput): Promise<import("axios").AxiosResponse<PostReviewResponse, any>>;
+export type PostReviewInput = z.input<typeof querySchema>;
+export declare function createReview({ path, content, startLine, rule, model, }: PostReviewInput): Promise<import("axios").AxiosResponse<PostReviewResponse, any>>;
 //# sourceMappingURL=_post.d.ts.map

actions/review-action/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@floe/review-action",
-  "version": "0.1.0-beta.6",
+  "version": "0.1.0-beta.7",
   "private": true,
   "scripts": {
     "build": "ncc build src/index.ts --out dist --source-map --license licenses.txt",

apps/api/CHANGELOG.md

@@ -1,5 +1,18 @@
 # @floe/api
 
+## 0.1.0-beta.7
+
+### Minor Changes
+
+- 5f84851: Add support for token usage and pro / basic models.
+
+### Patch Changes
+
+- Updated dependencies [5f84851]
+  - @floe/requests@0.1.0-beta.6
+  - @floe/lib@0.1.0-beta.4
+  - @floe/db@0.1.0-beta.2
+
 ## 0.1.0-beta.6
 
 ### Patch Changes

apps/api/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@floe/api",
-  "version": "0.1.0-beta.6",
+  "version": "0.1.0-beta.7",
   "private": true,
   "scripts": {
     "dev": "next dev -p 4000",

apps/api/src/lib/ai/index.ts

@@ -1,34 +1,119 @@
 import OpenAI from "openai";
+import { db } from "@floe/db";
+import { tokenUsage, subscription } from "@floe/db/models";
 import { Langfuse } from "langfuse";
+import * as Sentry from "@sentry/nextjs";
 import type { z, AnyZodObject } from "zod";
+import { HttpError } from "@floe/lib/http-error";
+import { getMonthYearTimestamp } from "@floe/lib/get-month-year";
 import { zParse } from "~/utils/z-parse";
 
 interface ProviderOptions {
   openai: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming;
 }
 
+async function persistTokenUsage({
+  workspaceId,
+  promptTokens,
+  completionTokens,
+  proModel,
+}: {
+  workspaceId: string;
+  promptTokens: number;
+  completionTokens: number;
+  proModel: boolean;
+}) {
+  const monthYear = getMonthYearTimestamp();
+  const promptKey = proModel ? "proPromptTokens" : "basePromptTokens";
+  const completionKey = proModel
+    ? "proCompletionTokens"
+    : "baseCompletionTokens";
+
+  await db.tokenUsage.upsert({
+    where: {
+      workspaceId_monthYear: {
+        workspaceId,
+        monthYear,
+      },
+    },
+    create: {
+      monthYear,
+      workspaceId,
+      [promptKey]: promptTokens,
+      [completionKey]: completionTokens,
+    },
+    update: {
+      [promptKey]: {
+        increment: promptTokens,
+      },
+      [completionKey]: {
+        increment: completionTokens,
+      },
+    },
+  });
+}
+
+const PRO_MODELS = ["gpt-4-1106-preview", "gpt-4"];
+
 /**
- * Generate a chat completion and log to Langfuse
+ * Generate a chat completion and log to DB and Langfuse.
+ * This function must be used to make all AI requests for tracking purposes.
  */
 export async function createCompletion<T extends AnyZodObject>({
   name,
-  userId,
   metadata,
   provider,
+  workspaceId,
   providerOptions,
   completionResponseSchema,
 }: {
   name: string;
-  userId: string;
   metadata: Record<string, string>;
   provider: keyof ProviderOptions;
+  workspaceId: string;
   providerOptions: ProviderOptions[typeof provider];
   completionResponseSchema: T;
 }) {
   if (!process.env.LANGFUSE_SECRET_KEY || !process.env.LANGFUSE_PUBLIC_KEY) {
     throw new Error("Missing LANGFUSE_SECRET_KEY or LANGFUSE_PUBLIC_KEY");
   }
 
+  const isUsingProModel = PRO_MODELS.includes(providerOptions.model);
+
+  const result = await Promise.all([
+    tokenUsage.findOne(workspaceId),
+    subscription.getTokenLimits(workspaceId),
+  ]);
+
+  const usedTokens = result[0];
+  const tokenLimits = result[1];
+
+  if (!usedTokens) {
+    throw new Error("Could not get token usage or token limits.");
+  }
+
+  const totalProTokens =
+    usedTokens.proCompletionTokens + usedTokens.proPromptTokens;
+  const totalBaseTokens =
+    usedTokens.baseCompletionTokens + usedTokens.basePromptTokens;
+
+  /**
+   * Check if the workspace has exceeded their token limit
+   */
+  if (isUsingProModel) {
+    if (totalProTokens >= tokenLimits.proTokenLimit) {
+      throw new HttpError({
+        statusCode: 402,
+        message: "You have exceeded your Pro token limit.",
+      });
+    }
+  } else if (totalBaseTokens >= tokenLimits.baseTokenLimit) {
+    throw new HttpError({
+      statusCode: 402,
+      message: "You have exceeded your Basic token limit.",
+    });
+  }
+
   const openai = new OpenAI({
     apiKey: process.env.OPENAI_API_KEY,
   });
@@ -40,7 +125,7 @@ export async function createCompletion<T extends AnyZodObject>({
 
   const trace = langfuse.trace({
     name,
-    userId,
+    userId: workspaceId,
     metadata: {
       env: process.env.NODE_ENV,
       ...metadata,
@@ -58,11 +143,21 @@ export async function createCompletion<T extends AnyZodObject>({
     input: providerOptions.messages,
   });
 
+  let chatCompletion: OpenAI.Chat.Completions.ChatCompletion;
+
   /**
    * Call the LLM
    * Eventually we may want to make calls to other providers here
    */
-  const chatCompletion = await openai.chat.completions.create(providerOptions);
+  try {
+    chatCompletion = await openai.chat.completions.create(providerOptions);
+  } catch (error) {
+    throw new HttpError({
+      statusCode: 500,
+      message: "Failed to get completion from LLM.",
+    });
+  }
+  const { usage } = chatCompletion;
 
   // End generation - sets endTime
   generation.end({
@@ -71,6 +166,22 @@ export async function createCompletion<T extends AnyZodObject>({
 
   await langfuse.shutdownAsync();
 
+  try {
+    await persistTokenUsage({
+      workspaceId,
+      promptTokens: usage?.prompt_tokens ?? 0,
+      completionTokens: usage?.completion_tokens ?? 0,
+      proModel: isUsingProModel,
+    });
+  } catch (error) {
+    /**
+     * We've gotten this far and have a generated response, so we don't want to
+     * throw an error here. Instead, we'll just log the error and continue.
+     */
+    Sentry.captureException(error);
+    console.error(error);
+  }
+
   /**
    * Safely parse the response from the LLM
    */