You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 

239 lines
5.7 KiB

import { dbGlobal } from "drizzle-pkg/lib/db";
import {
agentDocuments,
type AgentDocument,
type AgentDocumentType,
AgentDocumentTypes,
} from "drizzle-pkg/lib/schema/agent-document";
import { eq, desc, and, or, like, sql } from "drizzle-orm";
export { type AgentDocument, type AgentDocumentType, AgentDocumentTypes };
// ============ Types ============
export interface CreateDocumentInput {
userId: number;
sessionId?: string | null;
type?: AgentDocumentType;
title: string;
content: string;
tags?: string[] | null;
}
export interface UpdateDocumentInput {
title?: string;
content?: string;
tags?: string[] | null;
}
export interface ListDocumentsOptions {
userId: number;
type?: AgentDocumentType;
keyword?: string;
page?: number;
pageSize?: number;
}
export interface ListDocumentsResult {
items: Omit<AgentDocument, "content">[];
total: number;
page: number;
pageSize: number;
hasMore: boolean;
}
// ============ Helpers ============
const SUMMARY_MAX_CHARS = 200;
export function generateSummary(
content: string,
maxChars: number = SUMMARY_MAX_CHARS,
): string | null {
if (content.length <= maxChars) return null;
return content.slice(0, maxChars) + "...";
}
function serializeTags(tags: string[] | null | undefined): string | null {
if (!tags || tags.length === 0) return null;
return JSON.stringify(tags);
}
function deserializeTags(tags: string | null): string[] {
if (!tags) return [];
try {
const parsed = JSON.parse(tags);
return Array.isArray(parsed) ? parsed : [];
} catch {
return [];
}
}
// ============ CRUD ============
export async function createDocument(
input: CreateDocumentInput,
): Promise<AgentDocument> {
const summary = generateSummary(input.content);
const [inserted] = await dbGlobal
.insert(agentDocuments)
.values({
userId: input.userId,
sessionId: input.sessionId ?? null,
type: input.type ?? "document",
title: input.title,
content: input.content,
summary,
tags: serializeTags(input.tags),
})
.returning();
if (!inserted) throw new Error("Failed to create agent document");
return inserted;
}
export async function getDocumentById(
id: number,
userId: number,
): Promise<AgentDocument | null> {
const rows = await dbGlobal
.select()
.from(agentDocuments)
.where(and(eq(agentDocuments.id, id), eq(agentDocuments.userId, userId)))
.limit(1);
return rows[0] ?? null;
}
export async function listDocuments(
opts: ListDocumentsOptions,
): Promise<ListDocumentsResult> {
const page = opts.page ?? 1;
const pageSize = opts.pageSize ?? 20;
const conditions = [eq(agentDocuments.userId, opts.userId)];
if (opts.type) {
conditions.push(eq(agentDocuments.type, opts.type));
}
if (opts.keyword) {
const term = `%${opts.keyword}%`;
conditions.push(
or(
like(agentDocuments.title, term),
like(agentDocuments.content, term),
)!,
);
}
const where = and(...conditions);
const [rows, countResult] = await Promise.all([
dbGlobal
.select({
id: agentDocuments.id,
userId: agentDocuments.userId,
sessionId: agentDocuments.sessionId,
type: agentDocuments.type,
title: agentDocuments.title,
summary: agentDocuments.summary,
tags: agentDocuments.tags,
createdAt: agentDocuments.createdAt,
updatedAt: agentDocuments.updatedAt,
})
.from(agentDocuments)
.where(where)
.orderBy(desc(agentDocuments.updatedAt))
.limit(pageSize)
.offset((page - 1) * pageSize),
dbGlobal
.select({ count: sql<number>`count(*)` })
.from(agentDocuments)
.where(where),
]);
const total = countResult[0]?.count ?? 0;
return {
items: rows.map((r) => ({ ...r, tags: deserializeTags(r.tags) })),
total,
page,
pageSize,
hasMore: page * pageSize < total,
};
}
export async function updateDocument(
id: number,
userId: number,
input: UpdateDocumentInput,
): Promise<AgentDocument | null> {
const existing = await getDocumentById(id, userId);
if (!existing) return null;
const updateData: Record<string, unknown> = {};
if (input.title !== undefined) updateData.title = input.title;
if (input.content !== undefined) {
updateData.content = input.content;
updateData.summary = generateSummary(input.content);
}
if (input.tags !== undefined) {
updateData.tags = serializeTags(input.tags);
}
if (Object.keys(updateData).length === 0) {
return existing;
}
const [updated] = await dbGlobal
.update(agentDocuments)
.set(updateData)
.where(and(eq(agentDocuments.id, id), eq(agentDocuments.userId, userId)))
.returning();
return updated ?? null;
}
export async function deleteDocument(
id: number,
userId: number,
): Promise<boolean> {
const result = await dbGlobal
.delete(agentDocuments)
.where(and(eq(agentDocuments.id, id), eq(agentDocuments.userId, userId)))
.returning({ id: agentDocuments.id });
return result.length > 0;
}
// ============ Memory Context ============
export async function getMemoryContext(
userId: number,
): Promise<string | null> {
const rows = await dbGlobal
.select({
title: agentDocuments.title,
content: agentDocuments.content,
})
.from(agentDocuments)
.where(
and(
eq(agentDocuments.userId, userId),
eq(agentDocuments.type, "memory"),
),
)
.orderBy(desc(agentDocuments.updatedAt));
if (rows.length === 0) return null;
const parts = rows.map(
(r) => `## ${r.title}\n${r.content}`,
);
return `<memory>\n${parts.join("\n\n")}\n</memory>`;
}
// ============ Tags Helper ============
export { deserializeTags, serializeTags };