61 lines
2.1 KiB
TypeScript
61 lines
2.1 KiB
TypeScript
import { drizzle } from "drizzle-orm/node-postgres";
|
|
import { eq, asc } from "drizzle-orm";
|
|
import pg from "pg";
|
|
import * as schema from "./schema";
|
|
|
|
const { Pool } = pg;
|
|
|
|
if (!process.env.DATABASE_URL) {
|
|
throw new Error(
|
|
"DATABASE_URL must be set. Did you forget to provision a database?",
|
|
);
|
|
}
|
|
|
|
export const pool = new Pool({ connectionString: process.env.DATABASE_URL });
|
|
export const db = drizzle(pool, { schema });
|
|
|
|
export * from "./schema";
|
|
|
|
// ── Session history helper ──────────────────────────────────────────────────
|
|
|
|
/**
|
|
* Load the most recent conversation history for a session, capped by turn
|
|
* count and approximate token budget.
|
|
*
|
|
* @param sessionId Session to load history for
|
|
* @param maxTurns Maximum number of messages to return (default 8)
|
|
* @param maxTokens Approximate token budget — stops including older messages
|
|
* once cumulative token_count exceeds this (default 4000)
|
|
* @returns Array of { role, content } objects in chronological order
|
|
*/
|
|
export async function getSessionHistory(
|
|
sessionId: string,
|
|
maxTurns = 8,
|
|
maxTokens = 4000,
|
|
): Promise<Array<{ role: "user" | "assistant"; content: string }>> {
|
|
const rows = await db
|
|
.select({
|
|
role: schema.sessionMessages.role,
|
|
content: schema.sessionMessages.content,
|
|
tokenCount: schema.sessionMessages.tokenCount,
|
|
})
|
|
.from(schema.sessionMessages)
|
|
.where(eq(schema.sessionMessages.sessionId, sessionId))
|
|
.orderBy(asc(schema.sessionMessages.id));
|
|
|
|
// Take the most recent messages that fit within budget
|
|
const result: Array<{ role: "user" | "assistant"; content: string }> = [];
|
|
let totalTokens = 0;
|
|
|
|
// Walk from newest to oldest, then reverse
|
|
for (let i = rows.length - 1; i >= 0 && result.length < maxTurns; i--) {
|
|
const row = rows[i]!;
|
|
const tokens = row.tokenCount ?? Math.ceil(row.content.length / 4);
|
|
if (totalTokens + tokens > maxTokens && result.length > 0) break;
|
|
totalTokens += tokens;
|
|
result.push({ role: row.role as "user" | "assistant", content: row.content });
|
|
}
|
|
|
|
return result.reverse();
|
|
}
|