119 lines
4.1 KiB
PL/PgSQL
119 lines
4.1 KiB
PL/PgSQL
-- Initialize pgvector extension and research tables
|
|
CREATE EXTENSION IF NOT EXISTS vector;
|
|
CREATE EXTENSION IF NOT EXISTS pg_trgm;
|
|
|
|
-- Document storage
|
|
CREATE TABLE IF NOT EXISTS documents (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
filename TEXT NOT NULL,
|
|
doc_id TEXT,
|
|
mime_type TEXT,
|
|
file_path TEXT,
|
|
status TEXT DEFAULT 'pending',
|
|
page_count INTEGER DEFAULT 0,
|
|
full_text TEXT,
|
|
metadata JSONB DEFAULT '{}',
|
|
created_at TIMESTAMPTZ DEFAULT NOW(),
|
|
updated_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
|
|
-- Vector embeddings for similarity search
|
|
CREATE TABLE IF NOT EXISTS chunks (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
doc_id UUID REFERENCES documents(id) ON DELETE CASCADE,
|
|
content TEXT NOT NULL,
|
|
vector VECTOR(4096),
|
|
page_num INTEGER,
|
|
block_index INTEGER,
|
|
polygon JSONB,
|
|
chunk_type TEXT DEFAULT 'text',
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
|
|
-- Research memories/knowledge store
|
|
CREATE TABLE IF NOT EXISTS memories (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
session_id UUID REFERENCES documents(id),
|
|
content TEXT NOT NULL,
|
|
vector VECTOR(4096),
|
|
memory_type TEXT DEFAULT 'fact',
|
|
importance INTEGER DEFAULT 3,
|
|
source_doc_id UUID REFERENCES documents(id),
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
|
|
-- Research findings
|
|
CREATE TABLE IF NOT EXISTS findings (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
session_id UUID REFERENCES documents(id),
|
|
question TEXT,
|
|
answer TEXT,
|
|
summary TEXT,
|
|
vector VECTOR(4096),
|
|
relevant_chunks JSONB,
|
|
agent_name TEXT,
|
|
confidence FLOAT,
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
|
|
-- Research sessions
|
|
CREATE TABLE IF NOT EXISTS research_sessions (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
query TEXT NOT NULL,
|
|
status TEXT DEFAULT 'running',
|
|
documents JSONB DEFAULT '[]',
|
|
findings JSONB DEFAULT '[]',
|
|
created_at TIMESTAMPTZ DEFAULT NOW(),
|
|
completed_at TIMESTAMPTZ
|
|
);
|
|
|
|
-- Indexes
|
|
CREATE INDEX IF NOT EXISTS idx_chunks_vector ON chunks USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_vector ON memories USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
|
|
CREATE INDEX IF NOT EXISTS idx_findings_vector ON findings USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
|
|
CREATE INDEX IF NOT EXISTS idx_chunks_trgm ON chunks USING gin (content gin_trgm_ops);
|
|
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
|
|
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
|
|
CREATE INDEX IF NOT EXISTS idx_memories_session ON memories(session_id);
|
|
CREATE INDEX IF NOT EXISTS idx_findings_session ON findings(session_id);
|
|
CREATE INDEX IF NOT EXISTS idx_documents_status ON documents(status);
|
|
|
|
-- Pipeline intermediate stage storage
|
|
CREATE TABLE IF NOT EXISTS pipeline_stages (
|
|
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
|
|
session_id UUID REFERENCES research_sessions(id) ON DELETE CASCADE,
|
|
stage TEXT NOT NULL CHECK (stage IN ('triage','evidence','synthesis')),
|
|
output TEXT,
|
|
state JSONB DEFAULT '{}',
|
|
created_at TIMESTAMPTZ DEFAULT NOW()
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_pipeline_session ON pipeline_stages(session_id);
|
|
CREATE INDEX IF NOT EXISTS idx_pipeline_stage ON pipeline_stages(stage, session_id);
|
|
|
|
-- Create function for cosine similarity
|
|
CREATE OR REPLACE FUNCTION cosine_similarity(v1 VECTOR, v2 VECTOR)
|
|
RETURNS FLOAT AS $$
|
|
SELECT 1 - (v1 <-> v2) / 2;
|
|
$$ LANGUAGE SQL;
|
|
|
|
-- Create function for vector similarity search
|
|
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(4096), match_count INTEGER DEFAULT 10, min_score FLOAT DEFAULT 0)
|
|
RETURNS TABLE(id UUID, chunk_id UUID, content TEXT, doc_id UUID, page_num INT, similarity FLOAT)
|
|
LANGUAGE plpgsql
|
|
AS $$
|
|
BEGIN
|
|
RETURN QUERY
|
|
SELECT
|
|
m.id::TEXT,
|
|
m.id::TEXT,
|
|
m.content::TEXT,
|
|
m.doc_id::TEXT,
|
|
m.page_num::INT,
|
|
cosine_similarity(m.vector, query_embedding)::FLOAT
|
|
FROM chunks m
|
|
WHERE cosine_similarity(m.vector, query_embedding) >= min_score
|
|
ORDER BY m.vector <-> query_embedding
|
|
LIMIT match_count;
|
|
END;
|
|
$$;
|