Files
timmy-tower/lib/db/src/index.ts

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();
}