95 lines
3.1 KiB
PL/PgSQL
95 lines
3.1 KiB
PL/PgSQL
-- Enable pgvector extension
|
|
CREATE EXTENSION IF NOT EXISTS vector;
|
|
|
|
-- Knowledge Bases (Collections)
|
|
CREATE TABLE IF NOT EXISTS knowledge_bases (
|
|
id SERIAL PRIMARY KEY,
|
|
name VARCHAR(255) UNIQUE NOT NULL,
|
|
kb_type VARCHAR(50) NOT NULL DEFAULT 'generic',
|
|
description TEXT,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
tags TEXT,
|
|
metadata JSONB DEFAULT '{}'::jsonb
|
|
);
|
|
|
|
CREATE INDEX idx_kb_name ON knowledge_bases(name);
|
|
CREATE INDEX idx_kb_type ON knowledge_bases(kb_type);
|
|
CREATE INDEX idx_kb_created_at ON knowledge_bases(created_at);
|
|
|
|
-- Documents
|
|
CREATE TABLE IF NOT EXISTS documents (
|
|
id SERIAL PRIMARY KEY,
|
|
kb_id INTEGER NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
|
doc_id VARCHAR(255) NOT NULL,
|
|
content TEXT NOT NULL,
|
|
embedding vector(768), -- nomic-embed-text produces 768-dim vectors
|
|
metadata JSONB DEFAULT '{}'::jsonb,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(kb_id, doc_id)
|
|
);
|
|
|
|
CREATE INDEX idx_doc_kb_id ON documents(kb_id);
|
|
CREATE INDEX idx_doc_doc_id ON documents(doc_id);
|
|
CREATE INDEX idx_doc_metadata ON documents USING GIN(metadata);
|
|
CREATE INDEX idx_doc_embedding ON documents USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
|
|
|
|
-- Agent OS specific content types
|
|
CREATE TABLE IF NOT EXISTS agent_os_content (
|
|
id SERIAL PRIMARY KEY,
|
|
kb_id INTEGER NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
|
|
content_type VARCHAR(50) NOT NULL, -- standard, agent, workflow, command, product, spec
|
|
name VARCHAR(255) NOT NULL,
|
|
content TEXT NOT NULL,
|
|
file_path TEXT,
|
|
metadata JSONB DEFAULT '{}'::jsonb,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
UNIQUE(kb_id, content_type, name)
|
|
);
|
|
|
|
CREATE INDEX idx_aos_kb_id ON agent_os_content(kb_id);
|
|
CREATE INDEX idx_aos_content_type ON agent_os_content(content_type);
|
|
CREATE INDEX idx_aos_name ON agent_os_content(name);
|
|
CREATE INDEX idx_aos_metadata ON agent_os_content USING GIN(metadata);
|
|
|
|
-- Function to update updated_at timestamp
|
|
CREATE OR REPLACE FUNCTION update_updated_at_column()
|
|
RETURNS TRIGGER AS $$
|
|
BEGIN
|
|
NEW.updated_at = CURRENT_TIMESTAMP;
|
|
RETURN NEW;
|
|
END;
|
|
$$ language 'plpgsql';
|
|
|
|
-- Trigger to auto-update updated_at
|
|
CREATE TRIGGER update_kb_updated_at BEFORE UPDATE ON knowledge_bases
|
|
FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
-- Helper function for cosine similarity search
|
|
CREATE OR REPLACE FUNCTION search_documents(
|
|
p_kb_id INTEGER,
|
|
p_query_embedding vector(768),
|
|
p_limit INTEGER DEFAULT 10,
|
|
p_metadata_filter JSONB DEFAULT NULL
|
|
)
|
|
RETURNS TABLE (
|
|
doc_id VARCHAR,
|
|
content TEXT,
|
|
similarity FLOAT,
|
|
metadata JSONB
|
|
) AS $$
|
|
BEGIN
|
|
RETURN QUERY
|
|
SELECT
|
|
d.doc_id,
|
|
d.content,
|
|
1 - (d.embedding <=> p_query_embedding) AS similarity,
|
|
d.metadata
|
|
FROM documents d
|
|
WHERE d.kb_id = p_kb_id
|
|
AND (p_metadata_filter IS NULL OR d.metadata @> p_metadata_filter)
|
|
ORDER BY d.embedding <=> p_query_embedding
|
|
LIMIT p_limit;
|
|
END;
|
|
$$ LANGUAGE plpgsql;
|
|
|