feat: add brain_type column

backend/models/brain_entity.py

@@ -1,10 +1,16 @@
+from enum import Enum
 from typing import Optional
 from uuid import UUID
 
 from pydantic import BaseModel
 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:

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)

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;

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