From 9e7a5a239f091a17f303f7cfd332514d0935cee6 Mon Sep 17 00:00:00 2001
From: mamadoudicko <mamadoudicko100@gmail.com>
Date: Tue, 7 Nov 2023 11:51:13 +0100
Subject: [PATCH] feat: add brain_type column

---
 backend/models/brain_entity.py                |  7 ++++
 backend/models/databases/supabase/brains.py   |  3 +-
 ...00_add_field_brain_type_to_brain_table.sql | 34 +++++++++++++++++++
 scripts/tables.sql                            | 15 ++++++--
 4 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 scripts/20231106110000_add_field_brain_type_to_brain_table.sql

diff --git a/backend/models/brain_entity.py b/backend/models/brain_entity.py
index 528fe0ac836c..568bd8da2b31 100644
--- a/backend/models/brain_entity.py
+++ b/backend/models/brain_entity.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Optional
 from uuid import UUID
 
@@ -5,6 +6,11 @@
 from routes.authorizations.types import RoleEnum
 
 
+class BrainType(str, Enum):
+    DOC = "doc"
+    API = "api"
+
+
 class BrainEntity(BaseModel):
     brain_id: UUID
     name: str
@@ -16,6 +22,7 @@ class BrainEntity(BaseModel):
     status: Optional[str]
     prompt_id: Optional[UUID]
     last_update: str
+    brain_type: BrainType
 
     @property
     def id(self) -> UUID:
diff --git a/backend/models/databases/supabase/brains.py b/backend/models/databases/supabase/brains.py
index 564166bd6ced..338e8e7787b4 100644
--- a/backend/models/databases/supabase/brains.py
+++ b/backend/models/databases/supabase/brains.py
@@ -2,7 +2,7 @@
 from uuid import UUID
 
 from logger import get_logger
-from models.brain_entity import BrainEntity, MinimalBrainEntity, PublicBrain
+from models.brain_entity import BrainEntity, BrainType, MinimalBrainEntity, PublicBrain
 from models.databases.repository import Repository
 from pydantic import BaseModel
 
@@ -18,6 +18,7 @@ class CreateBrainProperties(BaseModel):
     max_tokens: Optional[int] = 256
     openai_api_key: Optional[str] = None
     prompt_id: Optional[UUID] = None
+    brain_type: Optional[BrainType] = BrainType.DOC
 
     def dict(self, *args, **kwargs):
         brain_dict = super().dict(*args, **kwargs)
diff --git a/scripts/20231106110000_add_field_brain_type_to_brain_table.sql b/scripts/20231106110000_add_field_brain_type_to_brain_table.sql
new file mode 100644
index 000000000000..12be1ff63804
--- /dev/null
+++ b/scripts/20231106110000_add_field_brain_type_to_brain_table.sql
@@ -0,0 +1,34 @@
+-- Check if the ENUM type 'brain_type' already exists
+DO $$
+BEGIN
+  IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'brain_type_enum') THEN
+    -- Create the ENUM type 'brain_type' if it doesn't exist
+    CREATE TYPE brain_type_enum AS ENUM ('doc', 'api');
+  END IF;
+END $$;
+
+-- Add a column 'brain_type' to the 'brains' table using the 'brain_type' ENUM type
+BEGIN;
+
+-- Add a column 'brain_type' to the 'brains' table as the 'brain_type' ENUM type with a default value 'doc'
+DO $$ 
+BEGIN 
+    IF NOT EXISTS (
+        SELECT 1
+        FROM information_schema.columns
+        WHERE table_name = 'brains'
+        AND column_name = 'brain_type'
+    ) THEN
+        ALTER TABLE brains ADD COLUMN brain_type brain_type_enum DEFAULT 'doc';
+    END IF;
+END $$;
+
+-- Insert a migration record if it doesn't exist
+INSERT INTO migrations (name) 
+SELECT '20231106110000_add_field_brain_type_to_brain_table'
+WHERE NOT EXISTS (
+    SELECT 1 FROM migrations WHERE name = '20231106110000_add_field_brain_type_to_brain_table'
+);
+
+-- Commit the changes
+COMMIT;
diff --git a/scripts/tables.sql b/scripts/tables.sql
index 3ade3918bc9c..4ae46c60a2cb 100644
--- a/scripts/tables.sql
+++ b/scripts/tables.sql
@@ -126,6 +126,14 @@ CREATE TABLE IF NOT EXISTS prompts (
     status VARCHAR(255) DEFAULT 'private'
 );
 
+DO $$ 
+BEGIN 
+IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'brain_type_enum') THEN
+  -- Create the ENUM type 'brain_type' if it doesn't exist
+  CREATE TYPE brain_type_enum AS ENUM ('doc', 'api');
+END IF;
+END $$;
+
 --- Create brains table
 CREATE TABLE IF NOT EXISTS brains (
   brain_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
@@ -137,7 +145,8 @@ CREATE TABLE IF NOT EXISTS brains (
   temperature FLOAT,
   openai_api_key TEXT,
   prompt_id UUID REFERENCES prompts(id),
-  last_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+  last_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+  brain_type brain_type_enum DEFAULT 'doc'
 );
 
 
@@ -390,9 +399,9 @@ CREATE POLICY "Access Quivr Storage 1jccrwz_2" ON storage.objects FOR UPDATE TO
 CREATE POLICY "Access Quivr Storage 1jccrwz_3" ON storage.objects FOR DELETE TO anon USING (bucket_id = 'quivr');
 
 INSERT INTO migrations (name) 
-SELECT '2023110607100000_add_api_brain_definition_table'
+SELECT '20231106110000_add_field_brain_type_to_brain_table'
 WHERE NOT EXISTS (
-    SELECT 1 FROM migrations WHERE name = '2023110607100000_add_api_brain_definition_table'
+    SELECT 1 FROM migrations WHERE name = '20231106110000_add_field_brain_type_to_brain_table'
 );