Skip to content

Simplify databunny api. #92

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

Merged
izadoesdev merged 3 commits into from
Aug 17, 2025
Merged

Simplify databunny api. #92

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c503505
fix: remove extra padding from overview tab
sbansal1999 Aug 5, 2025
6609e3e
refactor: update icon imports in databunny chat
sbansal1999 Aug 16, 2025
a51751a
feat: Simplify databunny api
sbansal1999 Aug 16, 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: 1 addition & 1 deletion apps/api/src/agent/handlers/chart-handler.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ export async function* handleChartResponse(
: getRandomMessage(noDataMessages),
data: {
hasVisualization: queryResult.data.length > 0,
chartType: parsedAiJson.chart_type,
chartType: parsedAiJson.chart_type || 'bar',
data: queryResult.data,
responseType: 'chart',
},
Expand Down
2 changes: 1 addition & 1 deletion apps/api/src/agent/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ export { processAssistantRequest } from './processor';
export {
AIPlanSchema,
AIResponseJsonSchema,
comprehensiveUnifiedPrompt,
comprehensiveSystemPrompt,
Comment thread
sbansal1999 marked this conversation as resolved.
} from './prompts/agent';
export { getAICompletion } from './utils/ai-client';
export { executeQuery } from './utils/query-executor';
Expand Down
43 changes: 11 additions & 32 deletions apps/api/src/agent/processor.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import type { User } from '@databuddy/auth';
import type { Website } from '@databuddy/shared';
import type { AssistantRequestType } from '../schemas';
import { handleChartResponse } from './handlers/chart-handler';
import { handleMetricResponse } from './handlers/metric-handler';
import { comprehensiveUnifiedPrompt } from './prompts/agent';
import { getAICompletion } from './utils/ai-client';
import type { StreamingUpdate } from './utils/stream-utils';
import { generateThinkingSteps } from './utils/stream-utils';
Expand All @@ -23,17 +23,8 @@ const unexpectedErrorMessages = [
'Something went a bit wonky on my end. Try asking me again?',
];

export interface AssistantRequest {
message: string;
website_id: string;
website_hostname: string;
model?: 'chat' | 'agent' | 'agent-max';
context?: {
previousMessages?: Array<{
role?: string;
content: string;
}>;
};
export interface AssistantRequest extends AssistantRequestType {
websiteHostname: string;
}

export interface AssistantContext {
Expand All @@ -50,34 +41,22 @@ export async function* processAssistantRequest(

try {
console.info('✅ [Assistant Processor] Input validated', {
message: request.message,
website_id: request.website_id,
website_hostname: request.website_hostname,
message: request.messages.at(-1),
websiteId: request.websiteId,
websiteHostname: request.websiteHostname,
});
Comment thread
sbansal1999 marked this conversation as resolved.

if (context.user?.role === 'ADMIN') {
context.debugInfo.validatedInput = {
message: request.message,
website_id: request.website_id,
website_hostname: request.website_hostname,
message: request.messages.at(-1),
websiteId: request.websiteId,
websiteHostname: request.websiteHostname,
};
}

const aiStart = Date.now();
const fullPrompt = comprehensiveUnifiedPrompt(
request.message,
request.website_id,
request.website_hostname,
'execute_chat',
request.context?.previousMessages?.map((msg) => ({
role: msg.role || 'user',
content: msg.content,
})),
undefined,
request.model
);

const aiResponse = await getAICompletion({ prompt: fullPrompt });

const aiResponse = await getAICompletion(request);
const aiTime = Date.now() - aiStart;

const parsedResponse = aiResponse.content;
Expand Down
25 changes: 3 additions & 22 deletions apps/api/src/agent/prompts/agent.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { z } from 'zod';
import type { AssistantRequestType } from '../../schemas/assistant-schemas';

export const AIResponseJsonSchema = z.object({
sql: z.string().nullable().optional(),
Expand Down Expand Up @@ -41,14 +42,11 @@ export const AIPlanSchema = z.object({
plan: z.array(z.string()),
});

export const comprehensiveUnifiedPrompt = (
userQuery: string,
export const comprehensiveSystemPrompt = (
websiteId: string,
websiteHostname: string,
mode: 'analysis_only' | 'execute_chat' | 'execute_agent_step',
previousMessages?: Array<{ role: string; content: string }>,
agentToolResult?: Record<string, unknown>,
_model?: 'chat' | 'agent' | 'agent-max'
_model?: AssistantRequestType['model']
) => `
Comment on lines +45 to 50

@coderabbitai coderabbitai Bot Aug 16, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Prompt still references <user_query> and includes invalid ClickHouse INTERVAL examples—please fix

  • The new flow no longer supplies <user_query>; the prompt should instruct the model to use “the latest user message in the provided messages” instead of <user_query>.
  • Several SQL examples use Postgres-like INTERVAL literals (e.g., INTERVAL '7' DAY). ClickHouse syntax is INTERVAL 7 DAY (numeric, no quotes). This will lead to invalid queries if the model copies these patterns. There’s also a QUALIFY usage in an example that ClickHouse doesn’t support—prefer subqueries or window + WHERE/HAVING alternatives.

Suggested changes:

  • Replace occurrences of “process the <user_query>” with “process the latest user message in the conversation”.
  • Fix INTERVAL literals: '7' DAY → 7 DAY, '30' DAY → 30 DAY, etc.
  • Replace QUALIFY examples with ClickHouse-compatible patterns (CTEs or nested subqueries with ROW_NUMBER() in a subquery and filter in outer SELECT).

Additionally, to enforce correctness at the schema level (and reduce reliance on defaults), consider switching to Zod discriminated unions so chart responses must include a chart type:

// Outside this hunk: refactor AIResponseJsonSchema to a discriminated union
const Base = z.object({
  sql: z.string().nullable().optional(),
  thinking_steps: z.array(z.string()).optional(),
});

const TextResp = Base.extend({
  response_type: z.literal('text'),
  text_response: z.string().nullable().optional(),
  chart_type: z.null().optional(),
});

const MetricResp = Base.extend({
  response_type: z.literal('metric'),
  metric_value: z.string(),
  metric_label: z.string().nullable().optional(),
  text_response: z.string().nullable().optional(),
  chart_type: z.null().optional(),
});

const ChartResp = Base.extend({
  response_type: z.literal('chart'),
  sql: z.string(), // require SQL for charts
  chart_type: z.enum([/* …all chart types… */]).default('bar'), // enforce presence with default
  text_response: z.string().nullable().optional(),
});

export const AIResponseJsonSchema = z.discriminatedUnion('response_type', [
  TextResp, MetricResp, ChartResp,
]);

This aligns with the PR note about discriminated unions and ensures the LLM output is structurally valid for each response type.

<persona>
You are Databunny, a world-class, specialized data analyst for the website ${websiteHostname}. You are precise, analytical, and secure. Your sole purpose is to help users understand their website's analytics data by providing insights, generating SQL queries, and creating visualizations.
Expand Down Expand Up @@ -168,25 +166,8 @@ You are Databunny, a world-class, specialized data analyst for the website ${web
<website_id>${websiteId}</website_id>
<website_hostname>${websiteHostname}</website_hostname>
<mode>${mode}</mode>
<user_query>${userQuery}</user_query>
<current_date_utc>${new Date().toISOString().split('T')[0]}</current_date_utc>
<current_timestamp_utc>${new Date().toISOString()}</current_timestamp_utc>
${
previousMessages && previousMessages.length > 0
? `
<conversation_history>
${previousMessages
.slice(-4)
.map(
(msg) =>
`<message role="${msg.role}">${msg.content?.substring(0, 200)}${msg.content?.length > 200 ? '...' : ''}</message>`
)
.join('\n')}
</conversation_history>
`
: ''
}
${agentToolResult ? `<agent_tool_result>${JSON.stringify(agentToolResult)}</agent_tool_result>` : ''}
</request_context>

<workflow_instructions>
Expand Down
27 changes: 18 additions & 9 deletions apps/api/src/agent/utils/ai-client.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,11 @@
import { createOpenRouter } from '@openrouter/ai-sdk-provider';
import { generateObject } from 'ai';
import type { z } from 'zod';
import { AIResponseJsonSchema } from '../prompts/agent';
import type { AssistantRequest } from '../processor';
import {
AIResponseJsonSchema,
comprehensiveSystemPrompt,
} from '../prompts/agent';

const openrouter = createOpenRouter({
apiKey: process.env.AI_API_KEY,
Expand All @@ -10,11 +14,6 @@ const openrouter = createOpenRouter({
const AI_MODEL = 'google/gemini-2.5-flash-lite-preview-06-17';
// const AI_MODEL = 'openrouter/horizon-beta';

interface AICompletionRequest {
prompt: string;
temperature?: number;
}

interface AICompletionResponse {
content: z.infer<typeof AIResponseJsonSchema>;
usage: {
Expand All @@ -25,15 +24,25 @@ interface AICompletionResponse {
}

export async function getAICompletion(
request: AICompletionRequest
request: AssistantRequest
): Promise<AICompletionResponse> {
const startTime = Date.now();

const systemPrompt = comprehensiveSystemPrompt(
request.websiteId,
request.websiteHostname,
'execute_chat',
request.model
);

try {
const chat = await generateObject({
model: openrouter.chat(AI_MODEL),
messages: [{ role: 'user', content: request.prompt }],
temperature: request.temperature ?? 0.1,
messages: [
Comment thread
sbansal1999 marked this conversation as resolved.
{ role: 'system', content: systemPrompt },
...request.messages,
],
temperature: 0.1,
Comment thread
sbansal1999 marked this conversation as resolved.
schema: AIResponseJsonSchema,
Comment thread
sbansal1999 marked this conversation as resolved.
});

Expand Down
11 changes: 5 additions & 6 deletions apps/api/src/routes/assistant.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,10 @@ export const assistant = new Elysia({ prefix: '/v1/assistant' })
.post(
'/stream',
async ({ body, user }: { body: AssistantRequestType; user: User }) => {
const { message, website_id, model, context } = body;
const { messages, websiteId, model } = body;

try {
const websiteValidation = await validateWebsite(website_id);
const websiteValidation = await validateWebsite(websiteId);

if (!websiteValidation.success) {
Comment thread
sbansal1999 marked this conversation as resolved.
Comment on lines 27 to 33

@coderabbitai coderabbitai Bot Aug 16, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Critical: Ensure websiteId in body matches authenticated header to prevent cross-website access

The route validates the website using the body.websiteId while auth likely relies on the X-Website-Id header. If these diverge and aren’t checked, a client could authenticate with one site and attempt to query another. Enforce equality.

Apply this diff:

-		async ({ body, user }: { body: AssistantRequestType; user: User }) => {
+		async ({
+			body,
+			user,
+			headers,
+		}: {
+			body: AssistantRequestType;
+			user: User;
+			headers: Record<string, string>;
+		}) => {
 			const { messages, websiteId, model } = body;
+			const headerWebsiteId = headers['x-website-id'];
+			if (headerWebsiteId && headerWebsiteId !== websiteId) {
+				return createStreamingResponse(
+					createErrorResponse('websiteId mismatch between header and body')
+				);
+			}

If websiteAuth already injects a canonical websiteId elsewhere, alternatively drop websiteId from the body and use the authenticated one everywhere.

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
async ({ body, user }: { body: AssistantRequestType; user: User }) => {
const { message, website_id, model, context } = body;
const { messages, websiteId, model } = body;
try {
const websiteValidation = await validateWebsite(website_id);
const websiteValidation = await validateWebsite(websiteId);
if (!websiteValidation.success) {
async ({
body,
user,
headers,
}: {
body: AssistantRequestType;
user: User;
headers: Record<string, string>;
}) => {
const { messages, websiteId, model } = body;
const headerWebsiteId = headers['x-website-id'];
if (headerWebsiteId && headerWebsiteId !== websiteId) {
return createStreamingResponse(
createErrorResponse('websiteId mismatch between header and body')
);
}
try {
const websiteValidation = await validateWebsite(websiteId);
if (!websiteValidation.success) {
// …
🤖 Prompt for AI Agents 
In apps/api/src/routes/assistant.ts around lines 27 to 33, the handler currently
validates websiteId from the request body but does not ensure it matches the
authenticated website id supplied via headers/middleware; enforce equality by
comparing body.websiteId to the authenticated website id (e.g., req.header or
user.websiteId injected by websiteAuth) and return a 403/validation error if
they differ, or remove/ignore websiteId from the body and always use the
canonical authenticated website id throughout the route (update subsequent calls
to validateWebsite/use websiteId accordingly).

@sbansal1999 sbansal1999 Aug 17, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We should add this, isn't it? @izadoesdev

@coderabbitai coderabbitai Bot Aug 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!

return createStreamingResponse(
Expand All @@ -45,11 +45,10 @@ export const assistant = new Elysia({ prefix: '/v1/assistant' })
}

const assistantRequest: AssistantRequest = {
message,
website_id,
website_hostname: website.domain,
messages,
websiteId,
websiteHostname: website.domain,
model: model || 'chat',
context,
};

const assistantContext: AssistantContext = {
Expand Down
48 changes: 19 additions & 29 deletions apps/api/src/schemas/assistant-schemas.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,33 +1,23 @@
import { t } from 'elysia';

export const AssistantRequestSchema = t.Object({
message: t.String(),
website_id: t.String(),
model: t.Optional(
t.Union([t.Literal('chat'), t.Literal('agent'), t.Literal('agent-max')])
),
context: t.Optional(
t.Object({
previousMessages: t.Optional(
t.Array(
t.Object({
role: t.Optional(t.String()),
content: t.String(),
})
)
export const AssistantRequestSchema = t.Object(
{
messages: t.Array(
t.Object(
{
role: t.Union([t.Literal('user'), t.Literal('assistant')]),
content: t.String(),
},
{ additionalProperties: false }
),
})
),
});
{ minItems: 1 }
),
websiteId: t.String(),
model: t.Optional(
t.Union([t.Literal('chat'), t.Literal('agent'), t.Literal('agent-max')])
),
},
{ additionalProperties: false }
);

export type AssistantRequestType = {
message: string;
website_id: string;
model?: 'chat' | 'agent' | 'agent-max';
context?: {
previousMessages?: Array<{
role?: string;
content: string;
}>;
};
};
export type AssistantRequestType = typeof AssistantRequestSchema.static;
2 changes: 1 addition & 1 deletion apps/dashboard/app/(main)/websites/[id]/_components/tabs/overview-tab.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -914,7 +914,7 @@ export function WebsiteOverviewTab({
);

return (
<div className="space-y-6 pt-6">
<div className="space-y-6">
<EventLimitIndicator />
<div className="grid grid-cols-1 gap-3 sm:grid-cols-2 sm:gap-4 md:grid-cols-3 lg:grid-cols-3 xl:grid-cols-6">
{[
Expand Down
Loading