mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-08-16 20:20:36 +00:00
cleanup(ene): read canonical SQL schema file for chat tables; remove embedded DDL string
This commit is contained in:
parent
e5705e5ad2
commit
152c7cd5b8
1 changed files with 75 additions and 154 deletions
|
|
@ -1,5 +1,6 @@
|
||||||
use crate::models::{ChatMessage, ChatSession, IngestionReceipt};
|
use crate::models::{ChatMessage, ChatSession, IngestionReceipt};
|
||||||
use anyhow::{Context, Result};
|
use anyhow::{Context, Result};
|
||||||
|
use std::path::Path;
|
||||||
use tokio_postgres::{Client, Config, NoTls};
|
use tokio_postgres::{Client, Config, NoTls};
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
|
|
@ -54,93 +55,13 @@ impl RdsSink {
|
||||||
|
|
||||||
/// Create tables and indexes if they do not exist.
|
/// Create tables and indexes if they do not exist.
|
||||||
async fn init_tables(&self) -> Result<()> {
|
async fn init_tables(&self) -> Result<()> {
|
||||||
// Try to enable pgvector — harmless if already enabled or unavailable.
|
// Canonical DDL lives in sql/ene_chat_schema.sql.
|
||||||
let _ = self
|
// Keeping it out of Rust source reduces schema drift across services.
|
||||||
.client
|
let ddl_path = Path::new(env!("CARGO_MANIFEST_DIR")).join("sql/ene_chat_schema.sql");
|
||||||
.batch_execute("CREATE EXTENSION IF NOT EXISTS vector")
|
let ddl = std::fs::read_to_string(&ddl_path)
|
||||||
.await;
|
.with_context(|| format!("read ENE chat schema DDL at {:?}", ddl_path))?;
|
||||||
|
|
||||||
let ddl = r#"
|
self.client.batch_execute(&ddl).await.context("init DDL")?;
|
||||||
CREATE TABLE IF NOT EXISTS ene.chat_sessions (
|
|
||||||
session_id TEXT PRIMARY KEY,
|
|
||||||
title TEXT,
|
|
||||||
agent TEXT,
|
|
||||||
model TEXT,
|
|
||||||
workspace_fingerprint TEXT,
|
|
||||||
workspace_root TEXT,
|
|
||||||
fork_parent_session_id TEXT,
|
|
||||||
compaction_count INTEGER NOT NULL DEFAULT 0,
|
|
||||||
compaction_summary TEXT,
|
|
||||||
message_count INTEGER NOT NULL DEFAULT 0,
|
|
||||||
token_input_total BIGINT NOT NULL DEFAULT 0,
|
|
||||||
token_output_total BIGINT NOT NULL DEFAULT 0,
|
|
||||||
created_at_ms BIGINT NOT NULL,
|
|
||||||
updated_at_ms BIGINT NOT NULL,
|
|
||||||
first_message_at_ms BIGINT,
|
|
||||||
last_message_at_ms BIGINT,
|
|
||||||
embedding vector(768),
|
|
||||||
meta JSONB NOT NULL DEFAULT '{}',
|
|
||||||
receipt TEXT,
|
|
||||||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS ene.chat_messages (
|
|
||||||
id BIGSERIAL PRIMARY KEY,
|
|
||||||
session_id TEXT NOT NULL REFERENCES ene.chat_sessions(session_id) ON DELETE CASCADE,
|
|
||||||
message_index INTEGER NOT NULL,
|
|
||||||
role TEXT NOT NULL,
|
|
||||||
blocks JSONB NOT NULL,
|
|
||||||
text_content TEXT,
|
|
||||||
token_input BIGINT NOT NULL DEFAULT 0,
|
|
||||||
token_output BIGINT NOT NULL DEFAULT 0,
|
|
||||||
token_cache_creation BIGINT NOT NULL DEFAULT 0,
|
|
||||||
token_cache_read BIGINT NOT NULL DEFAULT 0,
|
|
||||||
tool_calls JSONB NOT NULL DEFAULT '[]',
|
|
||||||
embedding vector(768),
|
|
||||||
receipt_hash TEXT,
|
|
||||||
created_at_ms BIGINT NOT NULL,
|
|
||||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
|
||||||
UNIQUE(session_id, message_index)
|
|
||||||
);
|
|
||||||
|
|
||||||
CREATE TABLE IF NOT EXISTS ene.ingestion_receipts (
|
|
||||||
receipt_id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
||||||
shim_name TEXT NOT NULL,
|
|
||||||
status TEXT NOT NULL DEFAULT 'pending',
|
|
||||||
sha256 TEXT NOT NULL,
|
|
||||||
record_count BIGINT NOT NULL DEFAULT 0,
|
|
||||||
source_path TEXT NOT NULL,
|
|
||||||
meta JSONB NOT NULL DEFAULT '{}',
|
|
||||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
||||||
);
|
|
||||||
|
|
||||||
-- Columns may already exist after a schema migration; ADD COLUMN IF NOT EXISTS
|
|
||||||
-- is idempotent on PostgreSQL ≥ 9.6.
|
|
||||||
ALTER TABLE ene.chat_sessions ADD COLUMN IF NOT EXISTS title TEXT;
|
|
||||||
ALTER TABLE ene.chat_sessions ADD COLUMN IF NOT EXISTS agent TEXT;
|
|
||||||
ALTER TABLE ene.chat_sessions ADD COLUMN IF NOT EXISTS model TEXT;
|
|
||||||
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_sessions_updated
|
|
||||||
ON ene.chat_sessions(updated_at_ms DESC);
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_sessions_workspace
|
|
||||||
ON ene.chat_sessions(workspace_fingerprint);
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_messages_session_order
|
|
||||||
ON ene.chat_messages(session_id, message_index);
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_messages_receipt
|
|
||||||
ON ene.chat_messages(receipt_hash);
|
|
||||||
"#;
|
|
||||||
// Full-text and JSONB indexes require non-empty text_content / tool_calls;
|
|
||||||
// create them separately so a pgvector-missing cluster still gets the rest.
|
|
||||||
let fts_ddl = r#"
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_messages_text_search
|
|
||||||
ON ene.chat_messages USING GIN(to_tsvector('english', COALESCE(text_content, '')));
|
|
||||||
CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
|
||||||
ON ene.chat_messages USING GIN(tool_calls jsonb_path_ops);
|
|
||||||
"#;
|
|
||||||
self.client.batch_execute(ddl).await.context("init DDL")?;
|
|
||||||
if let Err(e) = self.client.batch_execute(fts_ddl).await {
|
|
||||||
warn!("FTS index creation skipped (non-fatal): {}", e);
|
|
||||||
}
|
|
||||||
info!("RDS schema initialized");
|
info!("RDS schema initialized");
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
@ -159,32 +80,32 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let meta_json = serde_json::to_value(&s.meta)?;
|
let meta_json = serde_json::to_value(&s.meta)?;
|
||||||
self.client
|
self.client
|
||||||
.execute(
|
.execute(
|
||||||
"INSERT INTO ene.chat_sessions \
|
"INSERT INTO ene.chat_sessions \\
|
||||||
(session_id, title, agent, model, \
|
(session_id, title, agent, model, \\
|
||||||
workspace_fingerprint, workspace_root, fork_parent_session_id, \
|
workspace_fingerprint, workspace_root, fork_parent_session_id, \\
|
||||||
compaction_count, compaction_summary, message_count, \
|
compaction_count, compaction_summary, message_count, \\
|
||||||
token_input_total, token_output_total, \
|
token_input_total, token_output_total, \\
|
||||||
created_at_ms, updated_at_ms, first_message_at_ms, last_message_at_ms, \
|
created_at_ms, updated_at_ms, first_message_at_ms, last_message_at_ms, \\
|
||||||
embedding, meta, receipt, updated_at) \
|
embedding, meta, receipt, updated_at) \\
|
||||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17::vector,$18,$19,now()) \
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17::vector,$18,$19,now()) \\
|
||||||
ON CONFLICT (session_id) DO UPDATE SET \
|
ON CONFLICT (session_id) DO UPDATE SET \\
|
||||||
title = EXCLUDED.title, \
|
title = EXCLUDED.title, \\
|
||||||
agent = EXCLUDED.agent, \
|
agent = EXCLUDED.agent, \\
|
||||||
model = EXCLUDED.model, \
|
model = EXCLUDED.model, \\
|
||||||
workspace_fingerprint = EXCLUDED.workspace_fingerprint, \
|
workspace_fingerprint = EXCLUDED.workspace_fingerprint, \\
|
||||||
workspace_root = EXCLUDED.workspace_root, \
|
workspace_root = EXCLUDED.workspace_root, \\
|
||||||
fork_parent_session_id = EXCLUDED.fork_parent_session_id, \
|
fork_parent_session_id = EXCLUDED.fork_parent_session_id, \\
|
||||||
compaction_count = EXCLUDED.compaction_count, \
|
compaction_count = EXCLUDED.compaction_count, \\
|
||||||
compaction_summary = EXCLUDED.compaction_summary, \
|
compaction_summary = EXCLUDED.compaction_summary, \\
|
||||||
message_count = EXCLUDED.message_count, \
|
message_count = EXCLUDED.message_count, \\
|
||||||
token_input_total = EXCLUDED.token_input_total, \
|
token_input_total = EXCLUDED.token_input_total, \\
|
||||||
token_output_total = EXCLUDED.token_output_total, \
|
token_output_total = EXCLUDED.token_output_total, \\
|
||||||
updated_at_ms = EXCLUDED.updated_at_ms, \
|
updated_at_ms = EXCLUDED.updated_at_ms, \\
|
||||||
first_message_at_ms = EXCLUDED.first_message_at_ms, \
|
first_message_at_ms = EXCLUDED.first_message_at_ms, \\
|
||||||
last_message_at_ms = EXCLUDED.last_message_at_ms, \
|
last_message_at_ms = EXCLUDED.last_message_at_ms, \\
|
||||||
embedding = EXCLUDED.embedding, \
|
embedding = EXCLUDED.embedding, \\
|
||||||
meta = EXCLUDED.meta, \
|
meta = EXCLUDED.meta, \\
|
||||||
receipt = EXCLUDED.receipt, \
|
receipt = EXCLUDED.receipt, \\
|
||||||
updated_at = now()",
|
updated_at = now()",
|
||||||
&[
|
&[
|
||||||
&s.session_id,
|
&s.session_id,
|
||||||
|
|
@ -230,22 +151,22 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let tool_calls_json = serde_json::to_value(&msg.tool_calls)?;
|
let tool_calls_json = serde_json::to_value(&msg.tool_calls)?;
|
||||||
self.client
|
self.client
|
||||||
.execute(
|
.execute(
|
||||||
"INSERT INTO ene.chat_messages \
|
"INSERT INTO ene.chat_messages \\
|
||||||
(session_id, message_index, role, blocks, text_content, \
|
(session_id, message_index, role, blocks, text_content, \\
|
||||||
token_input, token_output, token_cache_creation, token_cache_read, \
|
token_input, token_output, token_cache_creation, token_cache_read, \\
|
||||||
tool_calls, embedding, receipt_hash, created_at_ms, created_at) \
|
tool_calls, embedding, receipt_hash, created_at_ms, created_at) \\
|
||||||
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::vector,$12,$13,now()) \
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11::vector,$12,$13,now()) \\
|
||||||
ON CONFLICT (session_id, message_index) DO UPDATE SET \
|
ON CONFLICT (session_id, message_index) DO UPDATE SET \\
|
||||||
role = EXCLUDED.role, \
|
role = EXCLUDED.role, \\
|
||||||
blocks = EXCLUDED.blocks, \
|
blocks = EXCLUDED.blocks, \\
|
||||||
text_content = EXCLUDED.text_content, \
|
text_content = EXCLUDED.text_content, \\
|
||||||
token_input = EXCLUDED.token_input, \
|
token_input = EXCLUDED.token_input, \\
|
||||||
token_output = EXCLUDED.token_output, \
|
token_output = EXCLUDED.token_output, \\
|
||||||
token_cache_creation = EXCLUDED.token_cache_creation, \
|
token_cache_creation = EXCLUDED.token_cache_creation, \\
|
||||||
token_cache_read = EXCLUDED.token_cache_read, \
|
token_cache_read = EXCLUDED.token_cache_read, \\
|
||||||
tool_calls = EXCLUDED.tool_calls, \
|
tool_calls = EXCLUDED.tool_calls, \\
|
||||||
embedding = EXCLUDED.embedding, \
|
embedding = EXCLUDED.embedding, \\
|
||||||
receipt_hash = EXCLUDED.receipt_hash, \
|
receipt_hash = EXCLUDED.receipt_hash, \\
|
||||||
created_at_ms = EXCLUDED.created_at_ms",
|
created_at_ms = EXCLUDED.created_at_ms",
|
||||||
&[
|
&[
|
||||||
&session_id,
|
&session_id,
|
||||||
|
|
@ -297,8 +218,8 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let meta_json = serde_json::to_value(&r.meta)?;
|
let meta_json = serde_json::to_value(&r.meta)?;
|
||||||
self.client
|
self.client
|
||||||
.execute(
|
.execute(
|
||||||
"INSERT INTO ene.ingestion_receipts \
|
"INSERT INTO ene.ingestion_receipts \\
|
||||||
(shim_name, status, sha256, record_count, source_path, meta) \
|
(shim_name, status, sha256, record_count, source_path, meta) \\
|
||||||
VALUES ($1, $2, $3, $4, $5, $6)",
|
VALUES ($1, $2, $3, $4, $5, $6)",
|
||||||
&[
|
&[
|
||||||
&r.shim_name,
|
&r.shim_name,
|
||||||
|
|
@ -323,16 +244,16 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let rows = self
|
let rows = self
|
||||||
.client
|
.client
|
||||||
.query(
|
.query(
|
||||||
"SELECT s.session_id, s.title, s.agent, s.model, \
|
"SELECT s.session_id, s.title, s.agent, s.model, \\
|
||||||
COUNT(m.id) AS match_count, \
|
COUNT(m.id) AS match_count, \\
|
||||||
MAX(ts_rank(to_tsvector('english', COALESCE(m.text_content,'')), \
|
MAX(ts_rank(to_tsvector('english', COALESCE(m.text_content,'')), \\
|
||||||
plainto_tsquery('english', $1))) AS rank \
|
plainto_tsquery('english', $1))) AS rank \\
|
||||||
FROM ene.chat_sessions s \
|
FROM ene.chat_sessions s \\
|
||||||
JOIN ene.chat_messages m ON m.session_id = s.session_id \
|
JOIN ene.chat_messages m ON m.session_id = s.session_id \\
|
||||||
WHERE to_tsvector('english', COALESCE(m.text_content,'')) \
|
WHERE to_tsvector('english', COALESCE(m.text_content,'')) \\
|
||||||
@@ plainto_tsquery('english', $1) \
|
@@ plainto_tsquery('english', $1) \\
|
||||||
GROUP BY s.session_id, s.title, s.agent, s.model \
|
GROUP BY s.session_id, s.title, s.agent, s.model \\
|
||||||
ORDER BY rank DESC \
|
ORDER BY rank DESC \\
|
||||||
LIMIT $2",
|
LIMIT $2",
|
||||||
&[&query, &limit],
|
&[&query, &limit],
|
||||||
)
|
)
|
||||||
|
|
@ -369,11 +290,11 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let rows = self
|
let rows = self
|
||||||
.client
|
.client
|
||||||
.query(
|
.query(
|
||||||
"SELECT session_id, title, agent, model, \
|
"SELECT session_id, title, agent, model, \\
|
||||||
1 - (embedding <=> $1::vector) AS similarity \
|
1 - (embedding <=> $1::vector) AS similarity \\
|
||||||
FROM ene.chat_sessions \
|
FROM ene.chat_sessions \\
|
||||||
WHERE embedding IS NOT NULL \
|
WHERE embedding IS NOT NULL \\
|
||||||
ORDER BY embedding <=> $1::vector \
|
ORDER BY embedding <=> $1::vector \\
|
||||||
LIMIT $2",
|
LIMIT $2",
|
||||||
&[&vec_str, &limit],
|
&[&vec_str, &limit],
|
||||||
)
|
)
|
||||||
|
|
@ -397,10 +318,10 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let rows = self
|
let rows = self
|
||||||
.client
|
.client
|
||||||
.query(
|
.query(
|
||||||
"SELECT session_id, title, agent, model, message_count, \
|
"SELECT session_id, title, agent, model, message_count, \\
|
||||||
token_input_total, token_output_total, created_at_ms, updated_at_ms \
|
token_input_total, token_output_total, created_at_ms, updated_at_ms \\
|
||||||
FROM ene.chat_sessions \
|
FROM ene.chat_sessions \\
|
||||||
ORDER BY updated_at_ms DESC \
|
ORDER BY updated_at_ms DESC \\
|
||||||
LIMIT $1",
|
LIMIT $1",
|
||||||
&[&limit],
|
&[&limit],
|
||||||
)
|
)
|
||||||
|
|
@ -428,8 +349,8 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let Some(sess) = self
|
let Some(sess) = self
|
||||||
.client
|
.client
|
||||||
.query_opt(
|
.query_opt(
|
||||||
"SELECT session_id, title, agent, model, message_count, \
|
"SELECT session_id, title, agent, model, message_count, \\
|
||||||
token_input_total, token_output_total, created_at_ms, updated_at_ms, meta \
|
token_input_total, token_output_total, created_at_ms, updated_at_ms, meta \\
|
||||||
FROM ene.chat_sessions WHERE session_id = $1",
|
FROM ene.chat_sessions WHERE session_id = $1",
|
||||||
&[&session_id],
|
&[&session_id],
|
||||||
)
|
)
|
||||||
|
|
@ -442,9 +363,9 @@ CREATE INDEX IF NOT EXISTS idx_chat_messages_tool_search
|
||||||
let msg_rows = self
|
let msg_rows = self
|
||||||
.client
|
.client
|
||||||
.query(
|
.query(
|
||||||
"SELECT message_index, role, blocks, text_content, \
|
"SELECT message_index, role, blocks, text_content, \\
|
||||||
token_input, token_output, tool_calls, created_at_ms \
|
token_input, token_output, tool_calls, created_at_ms \\
|
||||||
FROM ene.chat_messages \
|
FROM ene.chat_messages \\
|
||||||
WHERE session_id = $1 ORDER BY message_index",
|
WHERE session_id = $1 ORDER BY message_index",
|
||||||
&[&session_id],
|
&[&session_id],
|
||||||
)
|
)
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue