mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-09-11 14:36:35 +03:00
feat: add brain_type column to brain table (#1603)
Issue: https://github.com/StanGirard/quivr/pull/1601 Demo: https://github.com/StanGirard/quivr/assets/63923024/a7fe76a7-8809-4ec7-b1e6-d6fe36928a5c
This commit is contained in:
parent
9382762d01
commit
dd2bc1f6fc
@ -1,3 +1,4 @@
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
from uuid import UUID
|
||||
|
||||
@ -5,6 +6,11 @@ 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:
|
||||
|
@ -2,7 +2,7 @@ from typing import Optional
|
||||
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)
|
||||
|
@ -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;
|
@ -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'
|
||||
);
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user