From 9382762d01e10264c0ed0fb4580b4d6e0f272673 Mon Sep 17 00:00:00 2001 From: Mamadou DICKO <63923024+mamadoudicko@users.noreply.github.com> Date: Tue, 7 Nov 2023 14:03:50 +0100 Subject: [PATCH] feat: add api_brain_definition table (#1601) Issue: https://github.com/StanGirard/quivr/issues/1549 --- backend/models/ApiBrainDefinition.py | 12 ++++ backend/models/databases/repository.py | 12 ++++ backend/models/databases/supabase/__init__.py | 2 + .../supabase/api_brain_definition.py | 55 +++++++++++++++++++ .../models/databases/supabase/onboarding.py | 4 +- backend/models/databases/supabase/supabase.py | 3 + .../api_brain_definition/__init__.py | 0 .../add_api_brain_definition.py | 14 +++++ .../delete_api_brain_definition.py | 9 +++ .../get_api_brain_definition.py | 11 ++++ ...7100000_add_api_brain_definition_table.sql | 18 ++++++ scripts/tables.sql | 13 ++++- 12 files changed, 148 insertions(+), 5 deletions(-) create mode 100644 backend/models/ApiBrainDefinition.py create mode 100644 backend/models/databases/supabase/api_brain_definition.py create mode 100644 backend/repository/api_brain_definition/__init__.py create mode 100644 backend/repository/api_brain_definition/add_api_brain_definition.py create mode 100644 backend/repository/api_brain_definition/delete_api_brain_definition.py create mode 100644 backend/repository/api_brain_definition/get_api_brain_definition.py create mode 100644 scripts/2023110607100000_add_api_brain_definition_table.sql diff --git a/backend/models/ApiBrainDefinition.py b/backend/models/ApiBrainDefinition.py new file mode 100644 index 000000000..d1f3f51bc --- /dev/null +++ b/backend/models/ApiBrainDefinition.py @@ -0,0 +1,12 @@ +from uuid import UUID + +from pydantic import BaseModel + + +class ApiBrainDefinition(BaseModel): + brain_id: UUID + method: str + url: str + params: dict + search_params: dict + secrets: dict diff --git a/backend/models/databases/repository.py b/backend/models/databases/repository.py index 844314ac9..039f374af 100644 --- a/backend/models/databases/repository.py +++ b/backend/models/databases/repository.py @@ -279,3 +279,15 @@ class Repository(ABC): @abstractmethod def remove_onboarding_more_than_x_days(self, days: int): pass + + @abstractmethod + def get_api_brain_definition(self, brain_id: UUID): + pass + + @abstractmethod + def add_api_brain_definition(self, brain_id: UUID, config): + pass + + @abstractmethod + def delete_api_brain_definition(self, brain_id: UUID): + pass diff --git a/backend/models/databases/supabase/__init__.py b/backend/models/databases/supabase/__init__.py index 6a155075a..74990669a 100644 --- a/backend/models/databases/supabase/__init__.py +++ b/backend/models/databases/supabase/__init__.py @@ -9,3 +9,5 @@ from models.databases.supabase.onboarding import Onboarding from models.databases.supabase.prompts import Prompts from models.databases.supabase.user_usage import UserUsage from models.databases.supabase.vectors import Vector + +from models.databases.supabase.api_brain_definition import ApiBrainDefinitions diff --git a/backend/models/databases/supabase/api_brain_definition.py b/backend/models/databases/supabase/api_brain_definition.py new file mode 100644 index 000000000..dcef00967 --- /dev/null +++ b/backend/models/databases/supabase/api_brain_definition.py @@ -0,0 +1,55 @@ +from enum import Enum +from typing import Optional +from uuid import UUID + +from models.ApiBrainDefinition import ApiBrainDefinition +from models.databases.repository import Repository +from pydantic import BaseModel + + +class ApiMethod(str, Enum): + GET = "GET" + POST = "POST" + PUT = "PUT" + DELETE = "DELETE" + + +class CreateApiBrainDefinition(BaseModel): + brain_id: UUID + method: ApiMethod + url: str + params: dict + search_params: dict + secrets: dict + + +class ApiBrainDefinitions(Repository): + def __init__(self, supabase_client): + self.db = supabase_client + + def get_api_brain_definition(self, brain_id: UUID) -> Optional[ApiBrainDefinition]: + response = ( + self.db.table("api_brain_definition") + .select("*") + .filter("brain_id", "eq", brain_id) + .execute() + ) + if len(response.data) == 0: + return None + + return ApiBrainDefinition(**response.data[0]) + + def add_api_brain_definition( + self, brain_id: UUID, config: CreateApiBrainDefinition + ) -> Optional[ApiBrainDefinition]: + response = self.db.table("api_brain_definition").insert( + [{"brain_id": str(brain_id), **config.dict()}] + ) + if len(response.data) == 0: + return None + return ApiBrainDefinition(**response.data[0]) + + def delete_api_brain_definition(self, brain_id: UUID) -> None: + self.db.table("api_brain_definition").delete().filter( + "brain_id", "eq", str(brain_id) + ).execute() diff --git a/backend/models/databases/supabase/onboarding.py b/backend/models/databases/supabase/onboarding.py index 005e53fd6..3e49bf3a9 100644 --- a/backend/models/databases/supabase/onboarding.py +++ b/backend/models/databases/supabase/onboarding.py @@ -3,9 +3,7 @@ from typing import Optional from uuid import UUID from fastapi import HTTPException -from models.databases.repository import ( - Repository, # Assuming you have a repository class -) +from models.databases.repository import Repository from pydantic import BaseModel diff --git a/backend/models/databases/supabase/supabase.py b/backend/models/databases/supabase/supabase.py index e3efdf471..614b44ca7 100644 --- a/backend/models/databases/supabase/supabase.py +++ b/backend/models/databases/supabase/supabase.py @@ -1,5 +1,6 @@ from logger import get_logger from models.databases.supabase import ( + ApiBrainDefinitions, ApiKeyHandler, Brain, BrainSubscription, @@ -28,6 +29,7 @@ class SupabaseDB( Prompts, Notifications, Knowledges, + ApiBrainDefinitions, ): def __init__(self, supabase_client): self.db = supabase_client @@ -42,3 +44,4 @@ class SupabaseDB( Notifications.__init__(self, supabase_client) Knowledges.__init__(self, supabase_client) Onboarding.__init__(self, supabase_client) + ApiBrainDefinitions.__init__(self, supabase_client) diff --git a/backend/repository/api_brain_definition/__init__.py b/backend/repository/api_brain_definition/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/backend/repository/api_brain_definition/add_api_brain_definition.py b/backend/repository/api_brain_definition/add_api_brain_definition.py new file mode 100644 index 000000000..5bcc8e825 --- /dev/null +++ b/backend/repository/api_brain_definition/add_api_brain_definition.py @@ -0,0 +1,14 @@ +from uuid import UUID + +from models.databases.supabase.api_brain_definition import ( + CreateApiBrainDefinition, +) +from models.settings import get_supabase_db + + +def add_api_brain_definition( + brain_id: UUID, api_brain_configs: CreateApiBrainDefinition +) -> None: + supabase_db = get_supabase_db() + + supabase_db.add_api_brain_definition(brain_id, api_brain_configs) diff --git a/backend/repository/api_brain_definition/delete_api_brain_definition.py b/backend/repository/api_brain_definition/delete_api_brain_definition.py new file mode 100644 index 000000000..3ce505a28 --- /dev/null +++ b/backend/repository/api_brain_definition/delete_api_brain_definition.py @@ -0,0 +1,9 @@ +from uuid import UUID + +from models.settings import get_supabase_db + + +def delete_api_brain_definition(brain_id: UUID) -> None: + supabase_db = get_supabase_db() + + supabase_db.delete_api_brain_definition(brain_id) diff --git a/backend/repository/api_brain_definition/get_api_brain_definition.py b/backend/repository/api_brain_definition/get_api_brain_definition.py new file mode 100644 index 000000000..090b12cb9 --- /dev/null +++ b/backend/repository/api_brain_definition/get_api_brain_definition.py @@ -0,0 +1,11 @@ +from typing import Optional +from uuid import UUID + +from models import get_supabase_db +from models.ApiBrainDefinition import ApiBrainDefinition + + +def get_api_brain_definition(brain_id: UUID) -> Optional[ApiBrainDefinition]: + supabase_db = get_supabase_db() + + return supabase_db.get_api_brain_definition(brain_id) diff --git a/scripts/2023110607100000_add_api_brain_definition_table.sql b/scripts/2023110607100000_add_api_brain_definition_table.sql new file mode 100644 index 000000000..7c2739bcc --- /dev/null +++ b/scripts/2023110607100000_add_api_brain_definition_table.sql @@ -0,0 +1,18 @@ +-- Create the new table with 6 columns +CREATE TABLE IF NOT EXISTS api_brain_definition ( + brain_id UUID REFERENCES brains(brain_id), + method VARCHAR(255) CHECK (method IN ('GET', 'POST', 'PUT', 'DELETE')), + url VARCHAR(255), + params JSON, + search_params JSON, + secrets JSON +); + +-- Insert migration record if it doesn't exist +INSERT INTO migrations (name) +SELECT '2023110607100000_add_api_brain_definition_table' +WHERE NOT EXISTS ( + SELECT 1 FROM migrations WHERE name = '2023110607100000_add_api_brain_definition_table' +); + +COMMIT; \ No newline at end of file diff --git a/scripts/tables.sql b/scripts/tables.sql index 4dde7f6ea..3ade3918b 100644 --- a/scripts/tables.sql +++ b/scripts/tables.sql @@ -201,6 +201,15 @@ CREATE TABLE IF NOT EXISTS user_identity ( openai_api_key VARCHAR(255) ); +-- Create the new table with 6 columns +CREATE TABLE IF NOT EXISTS api_brain_definition ( + brain_id UUID REFERENCES brains(brain_id), + method VARCHAR(255) CHECK (method IN ('GET', 'POST', 'PUT', 'DELETE')), + url VARCHAR(255), + params JSON, + search_params JSON, + secrets JSON +); CREATE OR REPLACE FUNCTION public.get_user_email_by_user_id(user_id uuid) RETURNS TABLE (email text) @@ -381,9 +390,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 '20231023160000_copy_auth_users_to_public_users' +SELECT '2023110607100000_add_api_brain_definition_table' WHERE NOT EXISTS ( - SELECT 1 FROM migrations WHERE name = '20231023160000_copy_auth_users_to_public_users' + SELECT 1 FROM migrations WHERE name = '2023110607100000_add_api_brain_definition_table' );