Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chore: track ai usage #2271

Merged
merged 5 commits into from
Feb 11, 2025
Merged

chore: track ai usage #2271

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
2fecf19
chore: track ai usage
AyushAgrawal-A2 Feb 11, 2025
1f3fb9d
fix cache calculation
AyushAgrawal-A2 Feb 11, 2025
27c21ba
anthropic prompt caching
AyushAgrawal-A2 Feb 11, 2025
277f2b2
save to db, remove api rates
AyushAgrawal-A2 Feb 11, 2025
97e391e
Merge branch 'billing-v2' into ayush/ai_usage
AyushAgrawal-A2 Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions quadratic-api/prisma/migrations/20250211034001_optional_analytics/migration.sql
View file
Open in desktop

This file was deleted.

10 changes: 10 additions & 0 deletions ...grations/20250211191943_optional_analytics_and_token_usage_and_message_type/migration.sql
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
34 changes: 22 additions & 12 deletions quadratic-api/prisma/schema.prisma
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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])
Expand Down
12 changes: 6 additions & 6 deletions quadratic-api/src/ai/handler/anthropic.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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);

Expand All @@ -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) {
Expand All @@ -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);
Expand Down
20 changes: 10 additions & 10 deletions quadratic-api/src/ai/handler/bedrock.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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)) {
Expand All @@ -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) {
Expand All @@ -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);
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down
15 changes: 9 additions & 6 deletions quadratic-api/src/ai/handler/openai.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';

Expand All @@ -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);

Expand All @@ -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) {
Expand All @@ -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);
Expand Down
135 changes: 93 additions & 42 deletions quadratic-api/src/ai/helpers/anthropic.helper.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,20 +8,35 @@ import { aiToolsSpec } from 'quadratic-shared/ai/specs/aiToolsSpec';
import type {
AIMessagePrompt,
AIRequestBody,
AIUsage,
AnthropicModel,
BedrockAnthropicModel,
ParsedAIResponse,
} from 'quadratic-shared/typesAndSchemasAI';

export function getAnthropicApiArgs(args: Omit<AIRequestBody, 'chatId' | 'fileUuid' | 'source' | 'model'>): {
system: string | TextBlockParam[] | undefined;
system: TextBlockParam[] | undefined;
messages: MessageParam[];
tools: Tool[] | undefined;
tool_choice: ToolChoice | undefined;
} {
const { messages: chatMessages, useTools, toolName } = args;

const { systemMessages, promptMessages } = getSystemPromptMessages(chatMessages);
const system = systemMessages.join('\n\n');

// without prompt caching of system messages
const system: TextBlockParam[] = systemMessages.map((message) => ({
type: 'text' as const,
text: message,
}));

// with prompt caching of system messages
// const system: TextBlockParam[] = systemMessages.map((message, index) => ({
// type: 'text' as const,
// text: message,
// ...(index < 4 ? { cache_control: { type: 'ephemeral' } } : {}),
// }));

const messages: MessageParam[] = promptMessages.reduce<MessageParam[]>((acc, message) => {
if (message.role === 'assistant' && message.contextType === 'userPrompt' && message.toolCalls.length > 0) {
const anthropicMessages: MessageParam[] = [
Expand Down Expand Up @@ -120,7 +135,7 @@ export async function parseAnthropicStream(
chunks: Stream<Anthropic.Messages.RawMessageStreamEvent>,
response: Response,
model: AnthropicModel | BedrockAnthropicModel
) {
): Promise<ParsedAIResponse> {
const responseMessage: AIMessagePrompt = {
role: 'assistant',
content: '',
Expand All @@ -129,46 +144,75 @@ export async function parseAnthropicStream(
model,
};

const usage: AIUsage = {
inputTokens: 0,
outputTokens: 0,
cacheReadTokens: 0,
cacheWriteTokens: 0,
};

for await (const chunk of chunks) {
if (!response.writableEnded) {
if (chunk.type === 'content_block_start') {
if (chunk.content_block.type === 'text') {
responseMessage.content += chunk.content_block.text;
} else if (chunk.content_block.type === 'tool_use') {
const toolCalls = [...responseMessage.toolCalls];
const toolCall = {
id: chunk.content_block.id,
name: chunk.content_block.name,
arguments: '',
loading: true,
};
toolCalls.push(toolCall);
responseMessage.toolCalls = toolCalls;
}
} else if (chunk.type === 'content_block_delta') {
if (chunk.delta.type === 'text_delta') {
responseMessage.content += chunk.delta.text;
} else if (chunk.delta.type === 'input_json_delta') {
const toolCalls = [...responseMessage.toolCalls];
const toolCall = {
...(toolCalls.pop() ?? {
id: '',
name: '',
switch (chunk.type) {
case 'content_block_start':
if (chunk.content_block.type === 'text') {
responseMessage.content += chunk.content_block.text;
} else if (chunk.content_block.type === 'tool_use') {
const toolCalls = [...responseMessage.toolCalls];
const toolCall = {
id: chunk.content_block.id,
name: chunk.content_block.name,
arguments: '',
loading: true,
}),
};
toolCall.arguments += chunk.delta.partial_json;
toolCalls.push(toolCall);
responseMessage.toolCalls = toolCalls;
}
} else if (chunk.type === 'content_block_stop') {
const toolCalls = [...responseMessage.toolCalls];
const toolCall = toolCalls.pop();
if (toolCall) {
toolCalls.push({ ...toolCall, loading: false });
responseMessage.toolCalls = toolCalls;
}
};
toolCalls.push(toolCall);
responseMessage.toolCalls = toolCalls;
}
break;
case 'content_block_delta':
if (chunk.delta.type === 'text_delta') {
responseMessage.content += chunk.delta.text;
} else if (chunk.delta.type === 'input_json_delta') {
const toolCalls = [...responseMessage.toolCalls];
const toolCall = {
...(toolCalls.pop() ?? {
id: '',
name: '',
arguments: '',
loading: true,
}),
};
toolCall.arguments += chunk.delta.partial_json;
toolCalls.push(toolCall);
responseMessage.toolCalls = toolCalls;
}
break;
case 'content_block_stop':
{
const toolCalls = [...responseMessage.toolCalls];
const toolCall = toolCalls.pop();
if (toolCall) {
toolCalls.push({ ...toolCall, loading: false });
responseMessage.toolCalls = toolCalls;
}
}
break;
case 'message_start':
if (chunk.message.usage) {
usage.inputTokens = Math.max(usage.inputTokens, chunk.message.usage.input_tokens);
usage.outputTokens = Math.max(usage.outputTokens, chunk.message.usage.output_tokens);
usage.cacheReadTokens = Math.max(usage.cacheReadTokens, chunk.message.usage.cache_read_input_tokens ?? 0);
usage.cacheWriteTokens = Math.max(
usage.cacheWriteTokens,
chunk.message.usage.cache_creation_input_tokens ?? 0
);
}
break;
case 'message_delta':
if (chunk.usage) {
usage.outputTokens = Math.max(usage.outputTokens, chunk.usage.output_tokens);
}
break;
}

response.write(`data: ${JSON.stringify(responseMessage)}\n\n`);
Expand All @@ -187,14 +231,14 @@ export async function parseAnthropicStream(
response.end();
}

return responseMessage;
return { responseMessage, usage };
}

export function parseAnthropicResponse(
result: Anthropic.Messages.Message,
response: Response,
model: AnthropicModel | BedrockAnthropicModel
): AIMessagePrompt {
): ParsedAIResponse {
const responseMessage: AIMessagePrompt = {
role: 'assistant',
content: '',
Expand Down Expand Up @@ -233,5 +277,12 @@ export function parseAnthropicResponse(

response.json(responseMessage);

return responseMessage;
const usage: AIUsage = {
inputTokens: result.usage.input_tokens,
outputTokens: result.usage.output_tokens,
cacheReadTokens: result.usage.cache_read_input_tokens ?? 0,
cacheWriteTokens: result.usage.cache_creation_input_tokens ?? 0,
};

return { responseMessage, usage };
}
Loading
Loading