WIP: Reintroduce Billing

infra/aws-cloudformation/quadratic-preview.yml

@@ -139,6 +139,9 @@ Resources:
           ANTHROPIC_API_KEY=$(aws ssm get-parameter --name "/quadratic-development/ANTHROPIC_API_KEY" --with-decryption --query "Parameter.Value" --output text)
           EXA_API_KEY=$(aws ssm get-parameter --name "/quadratic-development/EXA_API_KEY" --with-decryption --query "Parameter.Value" --output text)
 
+          # stripe 
+          STRIPE_SECRET_KEY=$(aws ssm get-parameter --name "/quadratic-development/STRIPE_SECRET_KEY" --with-decryption --query "Parameter.Value" --output text)
+
           # aws ecr
           ECR_URL=${AWS::AccountId}.dkr.ecr.${AWS::Region}.amazonaws.com
           IMAGE_TAG=${ImageTag}

quadratic-api/prisma/migrations/20250211191943_optional_analytics_and_token_usage_and_message_type/migration.sql

@@ -0,0 +1,10 @@
+-- CreateEnum
+CREATE TYPE "AIChatMessageType" AS ENUM ('user_prompt', 'tool_result');
+
+-- AlterTable
+ALTER TABLE "AnalyticsAIChatMessage" ADD COLUMN     "cache_read_tokens" INTEGER,
+ADD COLUMN     "cache_write_tokens" INTEGER,
+ADD COLUMN     "input_tokens" INTEGER,
+ADD COLUMN     "message_type" "AIChatMessageType",
+ADD COLUMN     "output_tokens" INTEGER,
+ALTER COLUMN "s3_key" DROP NOT NULL;

quadratic-api/prisma/schema.prisma

@@ -285,19 +285,29 @@ enum AIChatSource {
     GetFileName  @map("get_file_name")
 }
 
+enum AIChatMessageType {
+    userPrompt @map("user_prompt")
+    toolResult @map("tool_result")
+}
+
 model AnalyticsAIChatMessage {
-    id            Int             @id @default(autoincrement())
-    chatId        Int             @map("chat_id")
-    chat          AnalyticsAIChat @relation(fields: [chatId], references: [id])
-    model         String
-    messageIndex  Int             @map("message_index")
-    s3Key         String          @map("s3_key")
-    like          Boolean?
-    undo          Boolean?
-    codeRunError  String?         @map("code_run_error")
-    responseError String?         @map("response_error")
-    createdDate   DateTime        @default(now()) @map("created_date")
-    updatedDate   DateTime        @default(now()) @map("updated_date")
+    id               Int                @id @default(autoincrement())
+    chatId           Int                @map("chat_id")
+    chat             AnalyticsAIChat    @relation(fields: [chatId], references: [id])
+    model            String
+    messageIndex     Int                @map("message_index")
+    messageType      AIChatMessageType? @map("message_type")
+    s3Key            String?            @map("s3_key")
+    like             Boolean?
+    undo             Boolean?
+    codeRunError     String?            @map("code_run_error")
+    responseError    String?            @map("response_error")
+    inputTokens      Int?               @map("input_tokens")
+    outputTokens     Int?               @map("output_tokens")
+    cacheReadTokens  Int?               @map("cache_read_tokens")
+    cacheWriteTokens Int?               @map("cache_write_tokens")
+    createdDate      DateTime           @default(now()) @map("created_date")
+    updatedDate      DateTime           @default(now()) @map("updated_date")
 
     @@unique([chatId, messageIndex], name: "chatId_messageIndex")
     @@index([chatId, messageIndex])

quadratic-api/src/ai/handler/anthropic.ts

@@ -1,7 +1,7 @@
 import Anthropic from '@anthropic-ai/sdk';
 import type { Response } from 'express';
 import { getModelOptions } from 'quadratic-shared/ai/helpers/model.helper';
-import type { AIMessagePrompt, AIRequestHelperArgs, AnthropicModel } from 'quadratic-shared/typesAndSchemasAI';
+import type { AIRequestHelperArgs, AnthropicModel, ParsedAIResponse } from 'quadratic-shared/typesAndSchemasAI';
 import { ANTHROPIC_API_KEY } from '../../env-vars';
 import { getAnthropicApiArgs, parseAnthropicResponse, parseAnthropicStream } from '../helpers/anthropic.helper';
 
@@ -13,7 +13,7 @@ export const handleAnthropicRequest = async (
   model: AnthropicModel,
   args: AIRequestHelperArgs,
   response: Response
-): Promise<AIMessagePrompt | undefined> => {
+): Promise<ParsedAIResponse | undefined> => {
   const { system, messages, tools, tool_choice } = getAnthropicApiArgs(args);
   const { stream, temperature, max_tokens } = getModelOptions(model, args);
 
@@ -34,8 +34,8 @@ export const handleAnthropicRequest = async (
       response.setHeader('Cache-Control', 'no-cache');
       response.setHeader('Connection', 'keep-alive');
 
-      const responseMessage = await parseAnthropicStream(chunks, response, model);
-      return responseMessage;
+      const parsedResponse = await parseAnthropicStream(chunks, response, model);
+      return parsedResponse;
     } catch (error: any) {
       if (!response.headersSent) {
         if (error instanceof Anthropic.APIError) {
@@ -62,8 +62,8 @@ export const handleAnthropicRequest = async (
         tool_choice,
       });
 
-      const responseMessage = parseAnthropicResponse(result, response, model);
-      return responseMessage;
+      const parsedResponse = parseAnthropicResponse(result, response, model);
+      return parsedResponse;
     } catch (error: any) {
       if (error instanceof Anthropic.APIError) {
         response.status(error.status ?? 400).json(error.message);

quadratic-api/src/ai/handler/bedrock.ts

@@ -3,7 +3,7 @@ import Anthropic from '@anthropic-ai/sdk';
 import { BedrockRuntimeClient, ConverseCommand, ConverseStreamCommand } from '@aws-sdk/client-bedrock-runtime';
 import { type Response } from 'express';
 import { getModelOptions, isBedrockAnthropicModel } from 'quadratic-shared/ai/helpers/model.helper';
-import type { AIMessagePrompt, AIRequestBody, BedrockModel } from 'quadratic-shared/typesAndSchemasAI';
+import type { AIRequestBody, BedrockModel, ParsedAIResponse } from 'quadratic-shared/typesAndSchemasAI';
 import { AWS_S3_ACCESS_KEY_ID, AWS_S3_REGION, AWS_S3_SECRET_ACCESS_KEY } from '../../env-vars';
 import { getAnthropicApiArgs, parseAnthropicResponse, parseAnthropicStream } from '../helpers/anthropic.helper';
 import { getBedrockApiArgs, parseBedrockResponse, parseBedrockStream } from '../helpers/bedrock.helper';
@@ -25,7 +25,7 @@ export const handleBedrockRequest = async (
   model: BedrockModel,
   args: Omit<AIRequestBody, 'chatId' | 'fileUuid' | 'source' | 'model'>,
   response: Response
-): Promise<AIMessagePrompt | undefined> => {
+): Promise<ParsedAIResponse | undefined> => {
   const { stream, temperature, max_tokens } = getModelOptions(model, args);
 
   if (isBedrockAnthropicModel(model)) {
@@ -47,8 +47,8 @@ export const handleBedrockRequest = async (
         response.setHeader('Cache-Control', 'no-cache');
         response.setHeader('Connection', 'keep-alive');
 
-        const responseMessage = await parseAnthropicStream(chunks, response, model);
-        return responseMessage;
+        const parsedResponse = await parseAnthropicStream(chunks, response, model);
+        return parsedResponse;
       } catch (error: any) {
         if (!response.headersSent) {
           if (error instanceof Anthropic.APIError) {
@@ -74,8 +74,8 @@ export const handleBedrockRequest = async (
           tool_choice,
         });
 
-        const responseMessage = parseAnthropicResponse(result, response, model);
-        return responseMessage;
+        const parsedResponse = parseAnthropicResponse(result, response, model);
+        return parsedResponse;
       } catch (error: any) {
         if (error instanceof Anthropic.APIError) {
           response.status(error.status ?? 400).json(error.message);
@@ -108,8 +108,8 @@ export const handleBedrockRequest = async (
         response.setHeader('Cache-Control', 'no-cache');
         response.setHeader('Connection', 'keep-alive');
 
-        const responseMessage = await parseBedrockStream(chunks, response, model);
-        return responseMessage;
+        const parsedResponse = await parseBedrockStream(chunks, response, model);
+        return parsedResponse;
       } catch (error: any) {
         if (!response.headersSent) {
           if (error.response) {
@@ -138,8 +138,8 @@ export const handleBedrockRequest = async (
         });
 
         const result = await bedrock.send(command);
-        const responseMessage = parseBedrockResponse(result.output, response, model);
-        return responseMessage;
+        const parsedResponse = parseBedrockResponse(result, response, model);
+        return parsedResponse;
       } catch (error: any) {
         if (error.response) {
           response.status(error.response.status).json(error.response.data);

quadratic-api/src/ai/handler/openai.ts

@@ -1,7 +1,7 @@
 import { type Response } from 'express';
 import OpenAI from 'openai';
 import { getModelOptions } from 'quadratic-shared/ai/helpers/model.helper';
-import type { AIMessagePrompt, AIRequestHelperArgs, OpenAIModel } from 'quadratic-shared/typesAndSchemasAI';
+import type { AIRequestHelperArgs, OpenAIModel, ParsedAIResponse } from 'quadratic-shared/typesAndSchemasAI';
 import { OPENAI_API_KEY } from '../../env-vars';
 import { getOpenAIApiArgs, parseOpenAIResponse, parseOpenAIStream } from '../helpers/openai.helper';
 
@@ -13,7 +13,7 @@ export const handleOpenAIRequest = async (
   model: OpenAIModel,
   args: AIRequestHelperArgs,
   response: Response
-): Promise<AIMessagePrompt | undefined> => {
+): Promise<ParsedAIResponse | undefined> => {
   const { messages, tools, tool_choice } = getOpenAIApiArgs(args);
   const { stream, temperature } = getModelOptions(model, args);
 
@@ -26,14 +26,17 @@ export const handleOpenAIRequest = async (
         stream: true,
         tools,
         tool_choice,
+        stream_options: {
+          include_usage: true,
+        },
       });
 
       response.setHeader('Content-Type', 'text/event-stream');
       response.setHeader('Cache-Control', 'no-cache');
       response.setHeader('Connection', 'keep-alive');
 
-      const responseMessage = await parseOpenAIStream(completion, response, model);
-      return responseMessage;
+      const parsedResponse = await parseOpenAIStream(completion, response, model);
+      return parsedResponse;
     } catch (error: any) {
       if (!response.headersSent) {
         if (error instanceof OpenAI.APIError) {
@@ -57,8 +60,8 @@ export const handleOpenAIRequest = async (
         tool_choice,
       });
 
-      const responseMessage = parseOpenAIResponse(result, response, model);
-      return responseMessage;
+      const parsedResponse = parseOpenAIResponse(result, response, model);
+      return parsedResponse;
     } catch (error: any) {
       if (error instanceof OpenAI.APIError) {
         response.status(error.status ?? 400).json(error.message);