@firtoz/chat-agent 1.0.0

package/src/chat-agent-drizzle.ts

@@ -0,0 +1,227 @@
+import { and, asc, eq, lt } from "drizzle-orm";
+import { migrate } from "drizzle-orm/durable-sqlite/migrator";
+import { createDb, type Database } from "./db/index";
+import {
+	messagesTable,
+	type NewMessage,
+	streamChunksTable,
+	streamMetadataTable,
+} from "./db/schema";
+import type { ChatMessage } from "./chat-messages";
+import { ChatAgentBase } from "./chat-agent-base";
+import migrations from "../drizzle/migrations.js";
+
+/**
+ * ChatAgent implementation using Drizzle ORM
+ *
+ * Uses Drizzle's type-safe query builder for database operations.
+ */
+export class DrizzleChatAgent<
+	Env extends Cloudflare.Env & {
+		OPENROUTER_API_KEY: string;
+	} = Cloudflare.Env & { OPENROUTER_API_KEY: string },
+> extends ChatAgentBase<Env> {
+	private db!: Database;
+
+	// ============================================================================
+	// Database Implementation - Drizzle ORM
+	// ============================================================================
+
+	protected dbInitialize(): void {
+		// Initialize Drizzle DB
+		this.db = createDb(this.ctx.storage);
+
+		// Run migrations
+		migrate(this.db, migrations);
+	}
+
+	protected dbLoadMessages(): ChatMessage[] {
+		const rows = this.db
+			.select()
+			.from(messagesTable)
+			.orderBy(asc(messagesTable.createdAt))
+			.all();
+
+		return rows
+			.map((row) => {
+				try {
+					return JSON.parse(row.messageJson) as ChatMessage;
+				} catch (err) {
+					console.error(`Failed to parse message ${row.id}:`, err);
+					return null;
+				}
+			})
+			.filter((msg): msg is ChatMessage => msg !== null);
+	}
+
+	protected dbSaveMessage(msg: ChatMessage): void {
+		const newMsg: NewMessage = {
+			id: msg.id,
+			role: msg.role,
+			messageJson: JSON.stringify(msg),
+			createdAt: new Date(msg.createdAt),
+		};
+
+		this.db
+			.insert(messagesTable)
+			.values(newMsg)
+			.onConflictDoUpdate({
+				target: messagesTable.id,
+				set: { messageJson: newMsg.messageJson },
+			})
+			.run();
+	}
+
+	protected dbClearAll(): void {
+		this.db.delete(messagesTable).run();
+		this.db.delete(streamChunksTable).run();
+		this.db.delete(streamMetadataTable).run();
+	}
+
+	protected dbFindActiveStream(): {
+		id: string;
+		messageId: string;
+		createdAt: Date;
+	} | null {
+		const activeStreams = this.db
+			.select()
+			.from(streamMetadataTable)
+			.where(eq(streamMetadataTable.status, "streaming"))
+			.orderBy(asc(streamMetadataTable.createdAt))
+			.limit(1)
+			.all();
+
+		if (activeStreams.length === 0) {
+			return null;
+		}
+
+		const stream = activeStreams[0];
+		return {
+			id: stream.id,
+			messageId: stream.messageId,
+			createdAt: stream.createdAt,
+		};
+	}
+
+	protected dbDeleteStreamWithChunks(streamId: string): void {
+		this.db
+			.delete(streamChunksTable)
+			.where(eq(streamChunksTable.streamId, streamId))
+			.run();
+		this.db
+			.delete(streamMetadataTable)
+			.where(eq(streamMetadataTable.id, streamId))
+			.run();
+	}
+
+	protected dbInsertStreamMetadata(streamId: string, messageId: string): void {
+		this.db
+			.insert(streamMetadataTable)
+			.values({
+				id: streamId,
+				messageId,
+				status: "streaming",
+				createdAt: new Date(),
+			})
+			.run();
+	}
+
+	protected dbUpdateStreamStatus(
+		streamId: string,
+		status: "completed" | "error",
+	): void {
+		this.db
+			.update(streamMetadataTable)
+			.set({ status, completedAt: new Date() })
+			.where(eq(streamMetadataTable.id, streamId))
+			.run();
+	}
+
+	protected dbDeleteOldCompletedStreams(cutoffMs: number): void {
+		const cutoff = new Date(cutoffMs);
+
+		// Delete old stream chunks for completed streams
+		const oldStreams = this.db
+			.select({ id: streamMetadataTable.id })
+			.from(streamMetadataTable)
+			.where(
+				and(
+					eq(streamMetadataTable.status, "completed"),
+					lt(streamMetadataTable.completedAt, cutoff),
+				),
+			)
+			.all();
+
+		for (const stream of oldStreams) {
+			this.db
+				.delete(streamChunksTable)
+				.where(eq(streamChunksTable.streamId, stream.id))
+				.run();
+		}
+
+		// Delete old stream metadata
+		this.db
+			.delete(streamMetadataTable)
+			.where(
+				and(
+					eq(streamMetadataTable.status, "completed"),
+					lt(streamMetadataTable.completedAt, cutoff),
+				),
+			)
+			.run();
+	}
+
+	protected dbFindMaxChunkIndex(streamId: string): number | null {
+		const lastChunk = this.db
+			.select({ maxIndex: streamChunksTable.chunkIndex })
+			.from(streamChunksTable)
+			.where(eq(streamChunksTable.streamId, streamId))
+			.orderBy(asc(streamChunksTable.chunkIndex))
+			.limit(1)
+			.all();
+
+		const firstChunk = lastChunk[0];
+		return firstChunk?.maxIndex != null ? firstChunk.maxIndex : null;
+	}
+
+	protected dbInsertChunks(
+		chunks: Array<{
+			id: string;
+			streamId: string;
+			content: string;
+			chunkIndex: number;
+		}>,
+	): void {
+		const now = new Date();
+		for (const chunk of chunks) {
+			this.db
+				.insert(streamChunksTable)
+				.values({
+					id: chunk.id,
+					streamId: chunk.streamId,
+					content: chunk.content,
+					chunkIndex: chunk.chunkIndex,
+					createdAt: now,
+				})
+				.run();
+		}
+	}
+
+	protected dbGetChunks(streamId: string): string[] {
+		const rows = this.db
+			.select()
+			.from(streamChunksTable)
+			.where(eq(streamChunksTable.streamId, streamId))
+			.orderBy(asc(streamChunksTable.chunkIndex))
+			.all();
+
+		return rows.map((r) => r.content);
+	}
+
+	protected dbDeleteChunks(streamId: string): void {
+		this.db
+			.delete(streamChunksTable)
+			.where(eq(streamChunksTable.streamId, streamId))
+			.run();
+	}
+}

package/src/chat-agent-sql.ts

@@ -0,0 +1,199 @@
+import type { ChatMessage } from "./chat-messages";
+import { ChatAgentBase } from "./chat-agent-base";
+
+/**
+ * ChatAgent implementation using raw SQL (like @cloudflare/ai-chat)
+ *
+ * Uses Agent's built-in `this.sql` template tag for database operations.
+ */
+export class SqlChatAgent<
+	Env extends Cloudflare.Env & {
+		OPENROUTER_API_KEY: string;
+	} = Cloudflare.Env & { OPENROUTER_API_KEY: string },
+> extends ChatAgentBase<Env> {
+	// ============================================================================
+	// Database Implementation - Raw SQL
+	// ============================================================================
+
+	protected dbInitialize(): void {
+		// Create tables for chat messages and resumable streaming
+		// Based on @cloudflare/ai-chat pattern from reference
+		this.sql`create table if not exists cf_ai_chat_agent_messages (
+			id text primary key,
+			message text not null,
+			created_at datetime default current_timestamp
+		)`;
+
+		this.sql`create table if not exists cf_ai_chat_stream_chunks (
+			id text primary key,
+			stream_id text not null,
+			body text not null,
+			chunk_index integer not null,
+			created_at integer not null
+		)`;
+
+		this.sql`create table if not exists cf_ai_chat_stream_metadata (
+			id text primary key,
+			request_id text not null,
+			status text not null,
+			created_at integer not null,
+			completed_at integer
+		)`;
+
+		this.sql`create index if not exists idx_stream_chunks_stream_id 
+			on cf_ai_chat_stream_chunks(stream_id, chunk_index)`;
+	}
+
+	protected dbLoadMessages(): ChatMessage[] {
+		const rows =
+			(this
+				.sql`select * from cf_ai_chat_agent_messages order by created_at` as Array<{
+				id: string;
+				message: string;
+			}>) || [];
+
+		return rows
+			.map((row) => {
+				try {
+					return JSON.parse(row.message) as ChatMessage;
+				} catch (err) {
+					console.error(`Failed to parse message ${row.id}:`, err);
+					return null;
+				}
+			})
+			.filter((msg): msg is ChatMessage => msg !== null);
+	}
+
+	protected dbSaveMessage(msg: ChatMessage): void {
+		const messageJson = JSON.stringify(msg);
+		this.sql`
+			insert into cf_ai_chat_agent_messages (id, message)
+			values (${msg.id}, ${messageJson})
+			on conflict(id) do update set message = excluded.message
+		`;
+	}
+
+	protected dbClearAll(): void {
+		this.sql`delete from cf_ai_chat_agent_messages`;
+		this.sql`delete from cf_ai_chat_stream_chunks`;
+		this.sql`delete from cf_ai_chat_stream_metadata`;
+	}
+
+	protected dbFindActiveStream(): {
+		id: string;
+		messageId: string;
+		createdAt: Date;
+	} | null {
+		const activeStreams = this.sql`
+			select * from cf_ai_chat_stream_metadata 
+			where status = 'streaming' 
+			order by created_at desc 
+			limit 1
+		` as Array<{
+			id: string;
+			request_id: string;
+			status: string;
+			created_at: number;
+			completed_at: number | null;
+		}>;
+
+		if (!activeStreams || activeStreams.length === 0) {
+			return null;
+		}
+
+		const stream = activeStreams[0];
+		return {
+			id: stream.id,
+			messageId: stream.request_id,
+			createdAt: new Date(stream.created_at),
+		};
+	}
+
+	protected dbDeleteStreamWithChunks(streamId: string): void {
+		this
+			.sql`delete from cf_ai_chat_stream_chunks where stream_id = ${streamId}`;
+		this.sql`delete from cf_ai_chat_stream_metadata where id = ${streamId}`;
+	}
+
+	protected dbInsertStreamMetadata(streamId: string, messageId: string): void {
+		const now = Date.now();
+		this.sql`
+			insert into cf_ai_chat_stream_metadata (id, request_id, status, created_at)
+			values (${streamId}, ${messageId}, 'streaming', ${now})
+		`;
+	}
+
+	protected dbUpdateStreamStatus(
+		streamId: string,
+		status: "completed" | "error",
+	): void {
+		const now = Date.now();
+		this.sql`
+			update cf_ai_chat_stream_metadata 
+			set status = ${status}, completed_at = ${now} 
+			where id = ${streamId}
+		`;
+	}
+
+	protected dbDeleteOldCompletedStreams(cutoffMs: number): void {
+		// Delete old stream chunks first
+		this.sql`
+			delete from cf_ai_chat_stream_chunks 
+			where stream_id in (
+				select id from cf_ai_chat_stream_metadata 
+				where status = 'completed' and completed_at < ${cutoffMs}
+			)
+		`;
+		// Then delete the metadata
+		this.sql`
+			delete from cf_ai_chat_stream_metadata 
+			where status = 'completed' and completed_at < ${cutoffMs}
+		`;
+	}
+
+	protected dbFindMaxChunkIndex(streamId: string): number | null {
+		const result = this.sql`
+			select max(chunk_index) as max_index 
+			from cf_ai_chat_stream_chunks 
+			where stream_id = ${streamId}
+		` as Array<{ max_index: number | null }>;
+
+		if (!result || result.length === 0 || result[0].max_index == null) {
+			return null;
+		}
+
+		return result[0].max_index;
+	}
+
+	protected dbInsertChunks(
+		chunks: Array<{
+			id: string;
+			streamId: string;
+			content: string;
+			chunkIndex: number;
+		}>,
+	): void {
+		const now = Date.now();
+		for (const chunk of chunks) {
+			this.sql`
+				insert into cf_ai_chat_stream_chunks (id, stream_id, body, chunk_index, created_at)
+				values (${chunk.id}, ${chunk.streamId}, ${chunk.content}, ${chunk.chunkIndex}, ${now})
+			`;
+		}
+	}
+
+	protected dbGetChunks(streamId: string): string[] {
+		const rows = this.sql`
+			select body from cf_ai_chat_stream_chunks 
+			where stream_id = ${streamId} 
+			order by chunk_index asc
+		` as Array<{ body: string }>;
+
+		return (rows || []).map((r) => r.body);
+	}
+
+	protected dbDeleteChunks(streamId: string): void {
+		this
+			.sql`delete from cf_ai_chat_stream_chunks where stream_id = ${streamId}`;
+	}
+}