Skip to content

Commit

Permalink
chat logging tests
Browse files Browse the repository at this point in the history
  • Loading branch information
@AyushAgrawal-A2
AyushAgrawal-A2 committed Jan 13, 2025
1 parent ab1343e commit ee97b7e
Show file tree
Hide file tree
Showing 9 changed files with 183 additions and 50 deletions.
174 changes: 153 additions & 21 deletions quadratic-api/src/routes/v0/ai.chat.POST.test.ts
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
import request from 'supertest';
import { app } from '../../app';
import dbClient from '../../dbClient';
import { getFile } from '../../middleware/getFile';
import { clearDb, createFile, createTeam, createUser } from '../../tests/testDataGenerator';

const auth0Id = 'user';

const payload = {
chatId: '00000000-0000-0000-0000-000000000000',
fileUuid: '11111111-1111-1111-1111-111111111111',
Expand All @@ -16,21 +20,6 @@ const payload = {
useQuadraticContext: false,
};

beforeAll(async () => {
const user = await createUser({ auth0Id: 'user' });
const team = await createTeam({ users: [{ userId: user.id, role: 'OWNER' }] });
await createFile({
data: {
uuid: payload.fileUuid,
name: 'Untitled',
ownerTeamId: team.id,
creatorUserId: user.id,
},
});
});

afterAll(clearDb);

jest.mock('@anthropic-ai/bedrock-sdk', () => ({
AnthropicBedrock: jest.fn().mockImplementation(() => ({
messages: {
Expand All @@ -53,18 +42,93 @@ jest.mock('@anthropic-ai/bedrock-sdk', () => ({
})),
}));

beforeAll(async () => {
const user = await createUser({ auth0Id });
const team = await createTeam({ users: [{ userId: user.id, role: 'OWNER' }] });
await createFile({
data: {
uuid: payload.fileUuid,
name: 'Untitled',
ownerTeamId: team.id,
creatorUserId: user.id,
},
});
});

afterAll(clearDb);

describe('POST /v0/ai/chat', () => {
describe('an unauthorized user', () => {
it('responds with a 401', async () => {
await request(app).post('/v0/ai/chat').send(payload).set('Authorization', `Bearer InvalidToken user`).expect(401);
describe('authentication', () => {
it('responds with a 401 when the token is invalid', async () => {
await request(app)
.post('/v0/ai/chat')
.send({ ...payload, chatId: '00000000-0000-0000-0000-000000000001' })
.set('Authorization', `Bearer InvalidToken user`)
.expect(401);
});

it('responds with model response when the token is valid', async () => {
await request(app)
.post('/v0/ai/chat')
.send({ ...payload, chatId: '00000000-0000-0000-0000-000000000002' })
.set('Authorization', `Bearer ValidToken user`)
.expect(200)
.expect(({ body }) => {
expect(body).toEqual({
role: 'assistant',
content: 'This is a mocked response from Claude',
contextType: 'userPrompt',
toolCalls: [
{
id: 'tool_123',
name: 'example_tool',
arguments: JSON.stringify({ param1: 'value1' }),
loading: false,
},
],
model: payload.model,
});
});

// wait for the chat to be saved
await new Promise((resolve) => setTimeout(resolve, 250));
});
});

describe('an authorized user', () => {
it('responds with a 200', async () => {
describe('Analytics AI Chat', () => {
beforeEach(async () => {
await dbClient.$transaction([
dbClient.analyticsAIChatMessage.deleteMany(),
dbClient.analyticsAIChat.deleteMany(),
]);
});

it('saves the chat in storage when analyticsAi is enabled', async () => {
const analyticsAIChatsBefore = await dbClient.analyticsAIChat.findMany();
expect(analyticsAIChatsBefore.length).toBe(0);

const user = await dbClient.user.findUnique({
where: {
auth0Id,
},
});
expect(user).not.toBeNull();
if (!user) {
throw new Error('User not found');
}

const {
file: { ownerTeam },
} = await getFile({ uuid: payload.fileUuid, userId: user.id });
expect(ownerTeam).not.toBeNull();
if (!ownerTeam) {
throw new Error('Owner team not found');
}
expect(ownerTeam.settingAnalyticsAi).toBe(true);

await request(app)
.post('/v0/ai/chat')
.send(payload)
.send({ ...payload, chatId: '00000000-0000-0000-0000-000000000003' })
.set('Authorization', `Bearer ValidToken user`)
.expect(200)
.expect(({ body }) => {
Expand All @@ -83,6 +147,74 @@ describe('POST /v0/ai/chat', () => {
model: payload.model,
});
});

// wait for the chat to be saved
await new Promise((resolve) => setTimeout(resolve, 250));

const analyticsAIChatsAfter = await dbClient.analyticsAIChat.findMany();
expect(analyticsAIChatsAfter.length).toBe(1);
});

it('does not save the chat in storage when analyticsAi is disabled', async () => {
const analyticsAIChatsBefore = await dbClient.analyticsAIChat.findMany();
expect(analyticsAIChatsBefore.length).toBe(0);

const user = await dbClient.user.findUnique({
where: {
auth0Id,
},
});
expect(user).not.toBeNull();
if (!user) {
throw new Error('User not found');
}

const {
file: { ownerTeam },
} = await getFile({ uuid: payload.fileUuid, userId: user.id });
expect(ownerTeam).not.toBeNull();
if (!ownerTeam) {
throw new Error('Owner team not found');
}
expect(ownerTeam.settingAnalyticsAi).toBe(true);

await request(app)
.patch(`/v0/teams/${ownerTeam.uuid}`)
.set('Authorization', `Bearer ValidToken user`)
.send({ settings: { analyticsAi: false } })

.expect(200)
.expect((res) => {
expect(res.body.settings.analyticsAi).toBe(false);
});

await request(app)
.post('/v0/ai/chat')
.set('Authorization', `Bearer ValidToken user`)
.send({ ...payload, chatId: '00000000-0000-0000-0000-000000000004' })
.expect(200)
.expect(({ body }) => {
expect(body).toEqual({
role: 'assistant',
content: 'This is a mocked response from Claude',
contextType: 'userPrompt',
toolCalls: [
{
id: 'tool_123',
name: 'example_tool',
arguments: JSON.stringify({ param1: 'value1' }),
loading: false,
},
],
model: payload.model,
});
});

// wait for the chat to be saved
await new Promise((resolve) => setTimeout(resolve, 250));

const analyticsAIChatsAfter = await dbClient.analyticsAIChat.findMany();
expect(analyticsAIChatsAfter.length).toBe(0);
});
});
});
21 changes: 10 additions & 11 deletions ...pi/src/routes/v0/ai.feedback.POST.test.ts → ...i/src/routes/v0/ai.feedback.PATCH.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,15 @@ beforeAll(async () => {
afterAll(clearDb);

describe('POST /v0/ai/feedback', () => {
describe('an unauthorized user', () => {
it('responds with a 401', async () => {
describe('authentication', () => {
it('responds with a 401 when the token is invalid', async () => {
await request(app)
.post('/v0/ai/feedback')
.patch('/v0/ai/feedback')
.set('Authorization', `Bearer InvalidToken user`)
.send({
...payload,
like: true,
})
.set('Authorization', `Bearer InvalidToken user`)
.expect(401);
});
});
Expand All @@ -51,13 +51,12 @@ describe('POST /v0/ai/feedback', () => {

// create a like
await request(app)
.post('/v0/ai/feedback')
.patch('/v0/ai/feedback')
.set('Authorization', `Bearer ValidToken user`)
.send({
...payload,
like: true,
})
.set('Accept', 'application/json')
.set('Authorization', `Bearer ValidToken user`)
.expect(200)
.expect(({ body }) => {
expect(body.message).toBe('Feedback received');
Expand All @@ -75,12 +74,12 @@ describe('POST /v0/ai/feedback', () => {

// set to dislike
await request(app)
.post('/v0/ai/feedback')
.patch('/v0/ai/feedback')
.set('Authorization', `Bearer ValidToken user`)
.send({
...payload,
like: false,
})
.set('Authorization', `Bearer ValidToken user`)
.expect(200)
.expect(({ body }) => {
expect(body.message).toBe('Feedback received');
Expand All @@ -98,12 +97,12 @@ describe('POST /v0/ai/feedback', () => {

// unset like
await request(app)
.post('/v0/ai/feedback')
.patch('/v0/ai/feedback')
.set('Authorization', `Bearer ValidToken user`)
.send({
...payload,
like: null,
})
.set('Authorization', `Bearer ValidToken user`)
.expect(200)
.expect(({ body }) => {
expect(body.message).toBe('Feedback received');
Expand Down
4 changes: 2 additions & 2 deletions ...tic-api/src/routes/v0/ai.feedback.POST.ts → ...ic-api/src/routes/v0/ai.feedback.PATCH.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,10 +10,10 @@ import type { RequestWithUser } from '../../types/Request';
export default [validateAccessToken, handler];

const schema = z.object({
body: ApiSchemas['/v0/ai/feedback.POST.request'],
body: ApiSchemas['/v0/ai/feedback.PATCH.request'],
});

async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/ai/feedback.POST.response']>) {
async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/ai/feedback.PATCH.response']>) {
const {
body: { chatId, messageIndex, like },
} = parseRequest(req, schema);
Expand Down
2 changes: 1 addition & 1 deletion quadratic-api/src/routes/v0/internal.user.$auth0Id.connections.$uuid.GET.ts
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { Request, Response } from 'express';
import type { Request, Response } from 'express';
import z from 'zod';
import dbClient from '../../dbClient';
import { validateM2MAuth } from '../../internal/validateM2MAuth';
Expand Down
8 changes: 4 additions & 4 deletions quadratic-api/src/routes/v0/teams.$uuid.PATCH.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/teams/:
} = req;
const {
userMakingRequest: { permissions },
team: { clientDataKv: exisitingClientDataKv },
team: { clientDataKv: existingClientDataKv },
} = await getTeam({ uuid, userId });

// Can they make the edits they’re trying to make?
Expand All @@ -40,8 +40,8 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/teams/:
throw new ApiError(403, 'User does not have permission to edit this team’s settings.');
}

// Validate exisiting data in the db
const validatedExisitingClientDataKv = validateClientDataKv(exisitingClientDataKv);
// Validate existing data in the db
const validatedExistingClientDataKv = validateClientDataKv(existingClientDataKv);

// Update the team with supplied data
const newTeam = await dbClient.team.update({
Expand All @@ -50,7 +50,7 @@ async function handler(req: RequestWithUser, res: Response<ApiTypes['/v0/teams/:
},
data: {
...(name ? { name } : {}),
...(clientDataKv ? { clientDataKv: { ...validatedExisitingClientDataKv, ...clientDataKv } } : {}),
...(clientDataKv ? { clientDataKv: { ...validatedExistingClientDataKv, ...clientDataKv } } : {}),
...(settings ? { settingAnalyticsAi: settings.analyticsAi } : {}),
},
});
Expand Down
7 changes: 4 additions & 3 deletions quadratic-api/src/routes/v0/teams.$uuid.invites.POST.ts
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import * as Sentry from '@sentry/node';
import { Response } from 'express';
import { ApiSchemas, ApiTypes } from 'quadratic-shared/typesAndSchemas';
import type { Response } from 'express';
import type { ApiTypes } from 'quadratic-shared/typesAndSchemas';
import { ApiSchemas } from 'quadratic-shared/typesAndSchemas';
import { z } from 'zod';
import { getUsers, getUsersByEmail } from '../../auth/auth';
import dbClient from '../../dbClient';
Expand All @@ -11,7 +12,7 @@ import { getTeam } from '../../middleware/getTeam';
import { userMiddleware } from '../../middleware/user';
import { validateAccessToken } from '../../middleware/validateAccessToken';
import { parseRequest } from '../../middleware/validateRequestSchema';
import { RequestWithUser } from '../../types/Request';
import type { RequestWithUser } from '../../types/Request';
import { ApiError } from '../../utils/ApiError';
import { firstRoleIsHigherThanSecond } from '../../utils/permissions';

Expand Down
7 changes: 4 additions & 3 deletions quadratic-api/src/routes/v0/teams.$uuid.users.$userId.PATCH.ts
View file
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
import { Response } from 'express';
import { ApiSchemas, ApiTypes } from 'quadratic-shared/typesAndSchemas';
import type { Response } from 'express';
import type { ApiTypes } from 'quadratic-shared/typesAndSchemas';
import { ApiSchemas } from 'quadratic-shared/typesAndSchemas';
import { z } from 'zod';
import dbClient from '../../dbClient';
import { getTeam } from '../../middleware/getTeam';
import { userMiddleware } from '../../middleware/user';
import { validateAccessToken } from '../../middleware/validateAccessToken';
import { parseRequest } from '../../middleware/validateRequestSchema';
import { RequestWithUser } from '../../types/Request';
import type { RequestWithUser } from '../../types/Request';
import { ApiError } from '../../utils/ApiError';
import { firstRoleIsHigherThanSecond } from '../../utils/permissions';

Expand Down
6 changes: 3 additions & 3 deletions quadratic-client/src/shared/api/apiClient.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -363,11 +363,11 @@ export const apiClient = {
},

ai: {
feedback(body: ApiTypes['/v0/ai/feedback.POST.request']) {
feedback(body: ApiTypes['/v0/ai/feedback.PATCH.request']) {
return fetchFromApi(
`/v0/ai/feedback`,
{ method: 'POST', body: JSON.stringify(body) },
ApiSchemas['/v0/ai/feedback.POST.response']
{ method: 'PATCH', body: JSON.stringify(body) },
ApiSchemas['/v0/ai/feedback.PATCH.response']
);
},
},
Expand Down
4 changes: 2 additions & 2 deletions quadratic-shared/typesAndSchemas.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -423,12 +423,12 @@ export const ApiSchemas = {
'/v0/ai/chat.POST.request': AIRequestBodySchema,
'/v0/ai/chat.POST.response': AIMessagePromptSchema,

'/v0/ai/feedback.POST.request': z.object({
'/v0/ai/feedback.PATCH.request': z.object({
chatId: z.string().uuid(),
messageIndex: z.number(),
like: z.boolean().nullable(),
}),
'/v0/ai/feedback.POST.response': z.object({
'/v0/ai/feedback.PATCH.response': z.object({
message: z.string(),
}),
};
Expand Down

0 comments on commit ee97b7e

Please sign in to comment.