mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-08-17 02:10:31 +03:00
feat(api-keys): added customization (#1802)
# Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate):
This commit is contained in:
parent
4c89812832
commit
31c1802084
@ -2,7 +2,6 @@ from secrets import token_hex
|
||||
from typing import List
|
||||
from uuid import uuid4
|
||||
|
||||
from asyncpg.exceptions import UniqueViolationError
|
||||
from fastapi import APIRouter, Depends
|
||||
from logger import get_logger
|
||||
from middlewares.auth import AuthBearer, get_current_user
|
||||
@ -38,23 +37,18 @@ async def create_api_key(current_user: UserIdentity = Depends(get_current_user))
|
||||
|
||||
new_key_id = uuid4()
|
||||
new_api_key = token_hex(16)
|
||||
api_key_inserted = False
|
||||
|
||||
while not api_key_inserted:
|
||||
try:
|
||||
# Attempt to insert new API key into database
|
||||
api_keys_repository.create_api_key(new_key_id, new_api_key, current_user.id)
|
||||
api_key_inserted = True
|
||||
|
||||
except UniqueViolationError:
|
||||
# Generate a new API key if the current one is already in use
|
||||
new_api_key = token_hex(16)
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating new API key: {e}")
|
||||
return {"api_key": "Error creating new API key."}
|
||||
try:
|
||||
# Attempt to insert new API key into database
|
||||
response = api_keys_repository.create_api_key(
|
||||
new_key_id, new_api_key, current_user.id, "api_key", 30, False
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating new API key: {e}")
|
||||
return {"api_key": "Error creating new API key."}
|
||||
logger.info(f"Created new API key for user {current_user.email}.")
|
||||
|
||||
return {"api_key": new_api_key, "key_id": str(new_key_id)}
|
||||
return response # type: ignore
|
||||
|
||||
|
||||
@api_key_router.delete(
|
||||
|
@ -4,3 +4,8 @@ from pydantic import BaseModel
|
||||
class ApiKey(BaseModel):
|
||||
api_key: str
|
||||
key_id: str
|
||||
days: int
|
||||
only_chat: bool
|
||||
name: str
|
||||
creation_time: str
|
||||
is_active: bool
|
||||
|
@ -7,7 +7,14 @@ from modules.api_key.entity.api_key import ApiKey
|
||||
|
||||
class ApiKeysInterface(ABC):
|
||||
@abstractmethod
|
||||
def create_api_key(self, new_key_id: UUID, new_api_key: str, user_id: UUID):
|
||||
def create_api_key(
|
||||
self,
|
||||
new_key_id: UUID,
|
||||
new_api_key: str,
|
||||
user_id: UUID,
|
||||
days: int,
|
||||
only_chat: bool,
|
||||
):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
|
@ -1,7 +1,9 @@
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
from uuid import UUID
|
||||
|
||||
from models.settings import get_supabase_client
|
||||
from modules.api_key.entity.api_key import ApiKey
|
||||
from modules.api_key.repository.api_key_interface import ApiKeysInterface
|
||||
|
||||
|
||||
@ -10,7 +12,9 @@ class ApiKeys(ApiKeysInterface):
|
||||
supabase_client = get_supabase_client()
|
||||
self.db = supabase_client # type: ignore
|
||||
|
||||
def create_api_key(self, new_key_id, new_api_key, user_id):
|
||||
def create_api_key(
|
||||
self, new_key_id, new_api_key, user_id, name, days=30, only_chat=False
|
||||
) -> Optional[ApiKey]:
|
||||
response = (
|
||||
self.db.table("api_keys")
|
||||
.insert(
|
||||
@ -19,6 +23,9 @@ class ApiKeys(ApiKeysInterface):
|
||||
"key_id": str(new_key_id),
|
||||
"user_id": str(user_id),
|
||||
"api_key": str(new_api_key),
|
||||
"name": str(name),
|
||||
"days": int(days),
|
||||
"only_chat": bool(only_chat),
|
||||
"creation_time": datetime.utcnow().strftime(
|
||||
"%Y-%m-%d %H:%M:%S"
|
||||
),
|
||||
@ -28,7 +35,9 @@ class ApiKeys(ApiKeysInterface):
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
return response
|
||||
if len(response.data) == 0:
|
||||
return None
|
||||
return ApiKey(**response.data[0])
|
||||
|
||||
def delete_api_key(self, key_id: str, user_id: UUID):
|
||||
return (
|
||||
@ -48,7 +57,7 @@ class ApiKeys(ApiKeysInterface):
|
||||
self.db.table("api_keys")
|
||||
.select("api_key", "creation_time")
|
||||
.filter("api_key", "eq", api_key)
|
||||
.filter("is_active", "eq", True)
|
||||
.filter("is_active", "eq", str(True))
|
||||
.execute()
|
||||
)
|
||||
return response
|
||||
|
33
scripts/20231203173900_new_api_key_format.sql
Normal file
33
scripts/20231203173900_new_api_key_format.sql
Normal file
@ -0,0 +1,33 @@
|
||||
DO $$
|
||||
BEGIN
|
||||
-- Add 'name' column if it does not exist
|
||||
IF NOT EXISTS (SELECT FROM pg_attribute WHERE attrelid = 'api_keys'::regclass AND attname = 'name') THEN
|
||||
ALTER TABLE api_keys ADD COLUMN name TEXT DEFAULT 'API_KEY';
|
||||
END IF;
|
||||
|
||||
-- Add 'days' column if it does not exist
|
||||
IF NOT EXISTS (SELECT FROM pg_attribute WHERE attrelid = 'api_keys'::regclass AND attname = 'days') THEN
|
||||
ALTER TABLE api_keys ADD COLUMN days INT DEFAULT 30;
|
||||
END IF;
|
||||
|
||||
-- Add 'only_chat' column if it does not exist
|
||||
IF NOT EXISTS (SELECT FROM pg_attribute WHERE attrelid = 'api_keys'::regclass AND attname = 'only_chat') THEN
|
||||
ALTER TABLE api_keys ADD COLUMN only_chat BOOLEAN DEFAULT false;
|
||||
END IF;
|
||||
|
||||
-- Optionally, update default values for existing rows if necessary
|
||||
-- UPDATE api_keys SET name = 'API_KEY' WHERE name IS NULL;
|
||||
-- UPDATE api_keys SET days = 30 WHERE days IS NULL;
|
||||
-- UPDATE api_keys SET only_chat = false WHERE only_chat IS NULL;
|
||||
|
||||
END $$;
|
||||
|
||||
|
||||
-- Update migrations table
|
||||
INSERT INTO migrations (name)
|
||||
SELECT '20231203173900_new_api_key_format'
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM migrations WHERE name = '20231203173900_new_api_key_format'
|
||||
);
|
||||
|
||||
COMMIT;
|
@ -1,225 +0,0 @@
|
||||
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
|
||||
|
||||
-- Create users table
|
||||
CREATE TABLE IF NOT EXISTS user_daily_usage(
|
||||
user_id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
|
||||
email TEXT,
|
||||
date TEXT,
|
||||
daily_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 functions for secrets in vault
|
||||
CREATE OR REPLACE FUNCTION insert_secret(name text, secret text)
|
||||
returns uuid
|
||||
language plpgsql
|
||||
security definer
|
||||
set search_path = public
|
||||
as $$
|
||||
begin
|
||||
return vault.create_secret(secret, name);
|
||||
end;
|
||||
$$;
|
||||
|
||||
|
||||
create or replace function read_secret(secret_name text)
|
||||
returns text
|
||||
language plpgsql
|
||||
security definer set search_path = public
|
||||
as $$
|
||||
declare
|
||||
secret text;
|
||||
begin
|
||||
select decrypted_secret from vault.decrypted_secrets where name =
|
||||
secret_name into secret;
|
||||
return secret;
|
||||
end;
|
||||
$$;
|
||||
|
||||
create or replace function delete_secret(secret_name text)
|
||||
returns text
|
||||
language plpgsql
|
||||
security definer set search_path = public
|
||||
as $$
|
||||
declare
|
||||
deleted_rows int;
|
||||
begin
|
||||
delete from vault.decrypted_secrets where name = secret_name;
|
||||
get diagnostics deleted_rows = row_count;
|
||||
if deleted_rows = 0 then
|
||||
return false;
|
||||
else
|
||||
return true;
|
||||
end if;
|
||||
end;
|
||||
$$;
|
||||
|
||||
|
||||
CREATE TABLE IF NOT EXISTS migrations (
|
||||
name VARCHAR(255) PRIMARY KEY,
|
||||
executed_at TIMESTAMPTZ DEFAULT current_timestamp
|
||||
);
|
||||
|
||||
INSERT INTO migrations (name)
|
||||
SELECT '20231107104700_setup_vault'
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM migrations WHERE name = '20231107104700_setup_vault'
|
||||
);
|
@ -112,6 +112,9 @@ $$;
|
||||
CREATE TABLE IF NOT EXISTS api_keys(
|
||||
key_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
|
||||
user_id UUID REFERENCES auth.users (id),
|
||||
name TEXT DEFAULT "API_KEY",
|
||||
days INT DEFAULT 30,
|
||||
only_chat BOOLEAN DEFAULT false,
|
||||
api_key TEXT UNIQUE,
|
||||
creation_time TIMESTAMP DEFAULT current_timestamp,
|
||||
deleted_time TIMESTAMP,
|
||||
@ -443,7 +446,7 @@ $$;
|
||||
|
||||
|
||||
INSERT INTO migrations (name)
|
||||
SELECT '20231128173900_remove_openai_api_key'
|
||||
SELECT '20231203173900_new_api_key_format'
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM migrations WHERE name = '20231128173900_remove_openai_api_key'
|
||||
SELECT 1 FROM migrations WHERE name = '20231203173900_new_api_key_format'
|
||||
);
|
||||
|
@ -112,6 +112,9 @@ $$;
|
||||
CREATE TABLE IF NOT EXISTS api_keys(
|
||||
key_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
|
||||
user_id UUID REFERENCES auth.users (id),
|
||||
name TEXT DEFAULT 'API_KEY',
|
||||
days INT DEFAULT 30,
|
||||
only_chat BOOLEAN DEFAULT false,
|
||||
api_key TEXT UNIQUE,
|
||||
creation_time TIMESTAMP DEFAULT current_timestamp,
|
||||
deleted_time TIMESTAMP,
|
||||
@ -441,9 +444,16 @@ begin
|
||||
end;
|
||||
$$;
|
||||
|
||||
create schema if not exists extensions;
|
||||
|
||||
create table if not exists
|
||||
extensions.wrappers_fdw_stats ();
|
||||
|
||||
grant all on extensions.wrappers_fdw_stats to service_role;
|
||||
|
||||
|
||||
INSERT INTO migrations (name)
|
||||
SELECT '20231128173900_remove_openai_api_key'
|
||||
SELECT '20231203173900_new_api_key_format'
|
||||
WHERE NOT EXISTS (
|
||||
SELECT 1 FROM migrations WHERE name = '20231128173900_remove_openai_api_key'
|
||||
SELECT 1 FROM migrations WHERE name = '20231203173900_new_api_key_format'
|
||||
);
|
||||
|
Loading…
Reference in New Issue
Block a user