forked from QuivrHQ/quivr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pg_tables.sql
179 lines (163 loc) · 4.53 KB
/
pg_tables.sql
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
-- Create users table
CREATE TABLE IF NOT EXISTS users(
user_id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
email TEXT,
date TEXT,
requests_count INT
);
-- Create chats table
CREATE TABLE IF NOT EXISTS chats(
chat_id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
user_id UUID REFERENCES users(user_id),
creation_time TIMESTAMP DEFAULT current_timestamp,
history JSONB,
chat_name TEXT
);
-- Create chat_history table
CREATE TABLE IF NOT EXISTS chat_history (
message_id UUID DEFAULT uuid_generate_v4(),
chat_id UUID REFERENCES chats(chat_id),
user_message TEXT,
assistant TEXT,
message_time TIMESTAMP DEFAULT current_timestamp,
PRIMARY KEY (chat_id, message_id)
);
-- Create vector extension
CREATE EXTENSION IF NOT EXISTS vector;
-- Create vectors table
CREATE TABLE IF NOT EXISTS vectors (
id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
content TEXT,
metadata JSONB,
embedding VECTOR(1536)
);
-- Create function to match vectors
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(1536), match_count INT, p_brain_id UUID)
RETURNS TABLE(
id UUID,
brain_id UUID,
content TEXT,
metadata JSONB,
embedding VECTOR(1536),
similarity FLOAT
) LANGUAGE plpgsql AS $$
#variable_conflict use_column
BEGIN
RETURN QUERY
SELECT
vectors.id,
brains_vectors.brain_id,
vectors.content,
vectors.metadata,
vectors.embedding,
1 - (vectors.embedding <=> query_embedding) AS similarity
FROM
vectors
INNER JOIN
brains_vectors ON vectors.id = brains_vectors.vector_id
WHERE brains_vectors.brain_id = p_brain_id
ORDER BY
vectors.embedding <=> query_embedding
LIMIT match_count;
END;
$$;
-- Create stats table
CREATE TABLE IF NOT EXISTS stats (
time TIMESTAMP,
chat BOOLEAN,
embedding BOOLEAN,
details TEXT,
metadata JSONB,
id INTEGER PRIMARY KEY GENERATED ALWAYS AS IDENTITY
);
-- Create summaries table
CREATE TABLE IF NOT EXISTS summaries (
id BIGSERIAL PRIMARY KEY,
document_id UUID REFERENCES vectors(id),
content TEXT,
metadata JSONB,
embedding VECTOR(1536)
);
-- Create function to match summaries
CREATE OR REPLACE FUNCTION match_summaries(query_embedding VECTOR(1536), match_count INT, match_threshold FLOAT)
RETURNS TABLE(
id BIGINT,
document_id UUID,
content TEXT,
metadata JSONB,
embedding VECTOR(1536),
similarity FLOAT
) LANGUAGE plpgsql AS $$
#variable_conflict use_column
BEGIN
RETURN QUERY
SELECT
id,
document_id,
content,
metadata,
embedding,
1 - (summaries.embedding <=> query_embedding) AS similarity
FROM
summaries
WHERE 1 - (summaries.embedding <=> query_embedding) > match_threshold
ORDER BY
summaries.embedding <=> query_embedding
LIMIT match_count;
END;
$$;
-- Create api_keys table
CREATE TABLE IF NOT EXISTS api_keys(
key_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
user_id UUID REFERENCES users(user_id),
api_key TEXT UNIQUE,
creation_time TIMESTAMP DEFAULT current_timestamp,
deleted_time TIMESTAMP,
is_active BOOLEAN DEFAULT true
);
-- Create brains table
CREATE TABLE IF NOT EXISTS brains (
brain_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
name TEXT,
status TEXT,
model TEXT,
max_tokens TEXT,
temperature FLOAT
);
-- Create brains X users table
CREATE TABLE IF NOT EXISTS brains_users (
brain_id UUID,
user_id UUID,
rights VARCHAR(255),
default_brain BOOLEAN DEFAULT false,
PRIMARY KEY (brain_id, user_id),
FOREIGN KEY (user_id) REFERENCES users(user_id),
FOREIGN KEY (brain_id) REFERENCES brains (brain_id)
);
-- Create brains X vectors table
CREATE TABLE IF NOT EXISTS brains_vectors (
brain_id UUID,
vector_id UUID,
file_sha1 TEXT,
PRIMARY KEY (brain_id, vector_id),
FOREIGN KEY (vector_id) REFERENCES vectors (id),
FOREIGN KEY (brain_id) REFERENCES brains (brain_id)
);
-- Create brains X vectors table
CREATE TABLE IF NOT EXISTS brain_subscription_invitations (
brain_id UUID,
email VARCHAR(255),
rights VARCHAR(255),
PRIMARY KEY (brain_id, email),
FOREIGN KEY (brain_id) REFERENCES brains (brain_id)
);
CREATE TABLE IF NOT EXISTS migrations (
name VARCHAR(255) PRIMARY KEY,
executed_at TIMESTAMPTZ DEFAULT current_timestamp
);
INSERT INTO migrations (name)
SELECT '202307111517031_change_vectors_id_type'
WHERE NOT EXISTS (
SELECT 1 FROM migrations WHERE name = '202307111517031_change_vectors_id_type'
);