feat: add api_brain_definition table (#1601)

Issue: https://github.com/StanGirard/quivr/issues/1549

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

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

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

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

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
 
 

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)

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)

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)

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)

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;

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