feat: store full databunny conversation (#106)

sbansal1999 · web-flow · commit 93ad8f849dbf · 2025-08-21T20:25:07.000+03:00
diff --git a/apps/api/src/agent/core/ai-service.ts b/apps/api/src/agent/core/ai-service.ts
@@ -38,7 +38,7 @@ export class AIService {
 			context.website.id,
 			context.website.domain,
 			'execute_chat',
-			context.model as 'chat' | 'agent' | 'agent-max'
+			context.model
 		);
 
 		try {
diff --git a/apps/api/src/agent/core/assistant-orchestrator.ts b/apps/api/src/agent/core/assistant-orchestrator.ts
@@ -1,7 +1,7 @@
 import type { User } from '@databuddy/auth';
-import type { Website } from '@databuddy/shared';
+import type { StreamingUpdate } from '@databuddy/shared';
+import { createId, type Website } from '@databuddy/shared';
 import type { AssistantRequestType } from '../../schemas';
-import type { StreamingUpdate } from '../utils/stream-utils';
 import { AIService } from './ai-service';
 import { AssistantSession, type SessionMetrics } from './assistant-session';
 import { ConversationRepository } from './conversation-repository';
@@ -30,6 +30,7 @@ export class AssistantOrchestrator {
 		try {
 			// Step 1: Generate AI response
 			const aiResponse = await this.aiService.generateResponse(session);
+			const aiMessageId = createId();
 
 			if (!aiResponse.content) {
 				session.log('AI response was empty');
@@ -42,12 +43,24 @@ export class AssistantOrchestrator {
 				];
 			}
 
+			const streamingUpdates: StreamingUpdate[] = [
+				{
+					type: 'metadata',
+					data: {
+						conversationId: session.getContext().conversationId,
+						messageId: aiMessageId,
+					},
+				},
+			];
+
 			// Step 2: Process the response into streaming updates
-			const streamingUpdates = await this.responseProcessor.process(
+			const aiResponseUpdates = await this.responseProcessor.process(
 				aiResponse.content,
 				session
 			);
 
+			streamingUpdates.push(...aiResponseUpdates);
+
 			// Step 3: Save to database (async, don't block response)
 			const finalResult = streamingUpdates.at(-1);
 			if (finalResult) {
@@ -57,6 +70,7 @@ export class AssistantOrchestrator {
 				this.saveConversationAsync(
 					session,
 					aiResponse.content,
+					aiMessageId,
 					finalResult,
 					metrics
 				);
@@ -85,13 +99,15 @@ export class AssistantOrchestrator {
 	private async saveConversationAsync(
 		session: AssistantSession,
 		aiResponse: AIResponseContent,
+		messageId: string,
 		finalResult: StreamingUpdate,
 		metrics: SessionMetrics
 	): Promise<void> {
 		try {
 			await this.conversationRepo.saveConversation(
 				session,
 				aiResponse,
+				messageId,
 				finalResult,
 				metrics
 			);
@@ -110,15 +126,21 @@ export class AssistantOrchestrator {
 		try {
 			const errorAIResponse = {
 				response_type: 'text' as const,
-				text_response: errorResponse.content,
+				text_response:
+					errorResponse.type === 'error'
+						? errorResponse.content
+						: 'Oops! Something unexpected happened. Mind trying that again?',
 				thinking_steps: [
 					`Error: ${originalError instanceof Error ? originalError.message : 'Unknown error'}`,
 				],
 			};
 
+			const messageId = createId('NANOID');
+
 			await this.conversationRepo.saveConversation(
 				session,
 				errorAIResponse,
+				messageId,
 				errorResponse,
 				metrics
 			);
diff --git a/apps/api/src/agent/core/assistant-session.ts b/apps/api/src/agent/core/assistant-session.ts
@@ -21,7 +21,7 @@ export interface SessionContext {
 	user: User;
 	website: Website;
 	conversationId: string;
-	model: string;
+	model: 'chat' | 'agent' | 'agent-max';
 }
 
 /**
diff --git a/apps/api/src/agent/core/conversation-repository.ts b/apps/api/src/agent/core/conversation-repository.ts
@@ -4,9 +4,9 @@ import {
 	assistantMessages,
 	db,
 } from '@databuddy/db';
+import type { StreamingUpdate } from '@databuddy/shared';
 import { createId } from '@databuddy/shared';
 import { eq } from 'drizzle-orm';
-import type { StreamingUpdate } from '../utils/stream-utils';
 import type { AssistantSession, SessionMetrics } from './assistant-session';
 import type { AIResponseContent } from './response-processor';
 
@@ -18,6 +18,7 @@ export class ConversationRepository {
 	async saveConversation(
 		session: AssistantSession,
 		aiResponse: AIResponseContent,
+		aiMessageId: string,
 		finalResult: StreamingUpdate,
 		metrics: SessionMetrics
 	): Promise<void> {
@@ -35,7 +36,6 @@ export class ConversationRepository {
 			messages.filter((m) => m.role === 'user').length === 1;
 
 		try {
-			const nowIso = new Date().toISOString();
 			const userMsg: AssistantMessageInput = {
 				id: createId('NANOID'),
 				conversationId: context.conversationId,
@@ -60,11 +60,10 @@ export class ConversationRepository {
 				totalTokens: null,
 				debugLogs: null,
 				metadata: null,
-				createdAt: nowIso,
 			};
 
 			const assistantMsg: AssistantMessageInput = {
-				id: createId('NANOID'),
+				id: aiMessageId,
 				conversationId: context.conversationId,
 				role: 'assistant',
 				content: finalResult.content,
@@ -77,17 +76,12 @@ export class ConversationRepository {
 				thinkingSteps: aiResponse.thinking_steps ?? null,
 				hasError: finalResult.type === 'error',
 				errorMessage: finalResult.type === 'error' ? finalResult.content : null,
-				upvotes: 0,
-				downvotes: 0,
-				feedbackComments: null,
 				aiResponseTime: metrics.aiResponseTime,
 				totalProcessingTime: metrics.totalProcessingTime,
 				promptTokens: metrics.tokenUsage.promptTokens,
 				completionTokens: metrics.tokenUsage.completionTokens,
 				totalTokens: metrics.tokenUsage.totalTokens,
 				debugLogs,
-				metadata: null,
-				createdAt: nowIso,
 			};
 
 			const conversationMessages: AssistantMessageInput[] = [
diff --git a/apps/api/src/agent/core/response-processor.ts b/apps/api/src/agent/core/response-processor.ts
@@ -2,7 +2,7 @@ import type { z } from 'zod';
 import { handleChartResponse } from '../handlers/chart-handler';
 import { handleMetricResponse } from '../handlers/metric-handler';
 import type { AIResponseJsonSchema } from '../prompts/agent';
-import type { StreamingUpdate } from '../utils/stream-utils';
+import type { StreamingUpdate } from '@databuddy/shared';
 import { generateThinkingSteps } from '../utils/stream-utils';
 import type { AssistantSession } from './assistant-session';
 
diff --git a/apps/api/src/agent/handlers/chart-handler.ts b/apps/api/src/agent/handlers/chart-handler.ts
@@ -1,8 +1,8 @@
+import type { StreamingUpdate } from '@databuddy/shared';
 import type { z } from 'zod';
 import type { AIResponseJsonSchema } from '../prompts/agent';
 import { executeQuery } from '../utils/query-executor';
 import { validateSQL } from '../utils/sql-validator';
-import type { StreamingUpdate } from '../utils/stream-utils';
 
 const getRandomMessage = (messages: string[]): string =>
 	messages[Math.floor(Math.random() * messages.length)] ||
diff --git a/apps/api/src/agent/handlers/metric-handler.ts b/apps/api/src/agent/handlers/metric-handler.ts
@@ -1,8 +1,8 @@
+import type { StreamingUpdate } from '@databuddy/shared';
 import type { z } from 'zod';
 import type { AIResponseJsonSchema } from '../prompts/agent';
 import { executeQuery } from '../utils/query-executor';
 import { validateSQL } from '../utils/sql-validator';
-import type { StreamingUpdate } from '../utils/stream-utils';
 
 export async function handleMetricResponse(
 	parsedAiJson: z.infer<typeof AIResponseJsonSchema>
diff --git a/apps/api/src/agent/processor.ts b/apps/api/src/agent/processor.ts
@@ -1,8 +1,7 @@
 import type { User } from '@databuddy/auth';
-import type { Website } from '@databuddy/shared';
+import type { StreamingUpdate, Website } from '@databuddy/shared';
 import type { AssistantRequestType } from '../schemas';
 import { AssistantOrchestrator } from './core/assistant-orchestrator';
-import type { StreamingUpdate } from './utils/stream-utils';
 
 /**
  * Clean, simple processor using object-oriented architecture
diff --git a/apps/api/src/agent/utils/stream-utils.ts b/apps/api/src/agent/utils/stream-utils.ts
@@ -1,13 +1,6 @@
-export interface StreamingUpdate {
-	type: 'thinking' | 'progress' | 'complete' | 'error';
-	content: string;
-	data?: Record<string, unknown>;
-	debugInfo?: Record<string, unknown>;
-}
+import type { StreamingUpdate } from '@databuddy/shared';
 
-export function createStreamingResponse(
-	updates: StreamingUpdate[]
-): Response {
+export function createStreamingResponse(updates: StreamingUpdate[]): Response {
 	const stream = new ReadableStream({
 		async start(controller) {
 			try {
@@ -43,9 +36,7 @@ function createThinkingStep(step: string): string {
 	return `🧠 ${step}`;
 }
 
-export function generateThinkingSteps(
-	steps: string[]
-): StreamingUpdate[] {
+export function generateThinkingSteps(steps: string[]): StreamingUpdate[] {
 	const updates: StreamingUpdate[] = [];
 	for (const step of steps) {
 		updates.push({ type: 'thinking', content: createThinkingStep(step) });
diff --git a/apps/api/src/routes/assistant.ts b/apps/api/src/routes/assistant.ts
@@ -1,10 +1,8 @@
 import { auth, type User, websitesApi } from '@databuddy/auth';
+import type { StreamingUpdate } from '@databuddy/shared';
 import { Elysia } from 'elysia';
 import { processAssistantRequest } from '../agent/processor';
-import {
-	createStreamingResponse,
-	type StreamingUpdate,
-} from '../agent/utils/stream-utils';
+import { createStreamingResponse } from '../agent/utils/stream-utils';
 import { validateWebsite } from '../lib/website-utils';
 import { AssistantRequestSchema, type AssistantRequestType } from '../schemas';
 
diff --git a/apps/dashboard/app/(main)/websites/[id]/assistant/components/chat-section.tsx b/apps/dashboard/app/(main)/websites/[id]/assistant/components/chat-section.tsx
@@ -27,14 +27,7 @@ import {
 import { Suggestion, Suggestions } from '@/components/ai-elements/suggestion';
 import { Skeleton } from '@/components/ui/skeleton';
 import { cn } from '@/lib/utils';
-import {
-	inputValueAtom,
-	isLoadingAtom,
-	isRateLimitedAtom,
-	messagesAtom,
-	modelAtom,
-	websiteDataAtom,
-} from '@/stores/jotai/assistantAtoms';
+import { modelAtom, websiteDataAtom } from '@/stores/jotai/assistantAtoms';
 import { useChat } from '../hooks/use-chat';
 import { MessageBubble } from './message-bubble';
 import { ModelSelector } from './model-selector';
@@ -88,15 +81,22 @@ const quickQuestions = [
 ];
 
 export default function ChatSection() {
-	const [messages] = useAtom(messagesAtom);
-	const [inputValue, setInputValue] = useAtom(inputValueAtom);
-	const [isLoading] = useAtom(isLoadingAtom);
-	const [isRateLimited] = useAtom(isRateLimitedAtom);
 	const [selectedModel] = useAtom(modelAtom);
 	const [websiteData] = useAtom(websiteDataAtom);
 
 	const inputRef = useRef<HTMLTextAreaElement>(null);
-	const { sendMessage, scrollToBottom, resetChat } = useChat();
+	const {
+		messages,
+		inputValue,
+		setInputValue,
+		isLoading,
+		isRateLimited,
+		sendMessage,
+		scrollToBottom,
+		resetChat,
+		handleUpVote,
+		handleDownVote,
+	} = useChat();
 
 	const hasMessages = messages.length > 1;
 
diff --git a/apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.ts b/apps/dashboard/app/(main)/websites/[id]/assistant/hooks/use-chat.ts
diff --git a/packages/shared/src/types/assistant.ts b/packages/shared/src/types/assistant.ts
diff --git a/packages/shared/src/types/index.ts b/packages/shared/src/types/index.ts

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ export interface SessionContext {`
`21`	`21`	`user: User;`
`22`	`22`	`website: Website;`
`23`	`23`	`conversationId: string;`
`24`		`- model: string;`
	`24`	`+ model: 'chat' \| 'agent' \| 'agent-max';`
`25`	`25`	`}`
`26`	`26`
`27`	`27`	`/**`