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):

backend/modules/api_key/controller/api_key_routes.py

@@ -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:
-        try:
+        # Attempt to insert new API key into database
-            # Attempt to insert new API key into database
+        response = api_keys_repository.create_api_key(
-            api_keys_repository.create_api_key(new_key_id, new_api_key, current_user.id)
+            new_key_id, new_api_key, current_user.id, "api_key", 30, False
-            api_key_inserted = True
+        )
-
+    except Exception as e:
-        except UniqueViolationError:
+        logger.error(f"Error creating new API key: {e}")
-            # Generate a new API key if the current one is already in use
+        return {"api_key": "Error creating new API key."}
-            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."}
     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(

backend/modules/api_key/entity/api_key.py

@@ -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

backend/modules/api_key/repository/api_key_interface.py

@@ -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

backend/modules/api_key/repository/api_keys.py

@@ -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

scripts/20231203173900_new_api_key_format.sql

@@ -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;

scripts/pg_tables.sql

@@ -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'
-);

scripts/tables-ollama.sql

@@ -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'
 );

scripts/tables.sql

@@ -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'
 );