From fb3aae27f1ee4d49d714eea3a569eccc247dacf8 Mon Sep 17 00:00:00 2001 From: Damien Mourot Date: Sun, 28 Apr 2024 22:27:50 -1000 Subject: [PATCH] feat(backend): use SQLAlchemy instead od supabase API (#2516) # Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate): --- backend/models/settings.py | 10 ++++ backend/modules/brain/entity/brain_entity.py | 3 +- backend/modules/brain/repository/brains.py | 52 +++++++++++--------- 3 files changed, 41 insertions(+), 24 deletions(-) diff --git a/backend/models/settings.py b/backend/models/settings.py index 33c487634..b4635e6b9 100644 --- a/backend/models/settings.py +++ b/backend/models/settings.py @@ -3,6 +3,7 @@ from uuid import UUID from langchain.embeddings.ollama import OllamaEmbeddings from langchain_openai import OpenAIEmbeddings +from sqlalchemy import Engine, create_engine from logger import get_logger from models.databases.supabase.supabase import SupabaseDB from posthog import Posthog @@ -114,6 +115,7 @@ class BrainSettings(BaseSettings): ollama_api_base_url: str = None langfuse_public_key: str = None langfuse_secret_key: str = None + pg_database_url: str = None class ResendSettings(BaseSettings): @@ -124,7 +126,15 @@ class ResendSettings(BaseSettings): # Global variables to store the Supabase client and database instances _supabase_client: Optional[Client] = None _supabase_db: Optional[SupabaseDB] = None +_db_engine: Optional[Engine] = None +def get_pg_database_engine(): + global _db_engine + if _db_engine is None: + logger.info("Creating Postgres DB engine") + settings = BrainSettings() # pyright: ignore reportPrivateUsage=none + _db_engine = create_engine(settings.pg_database_url) + return _db_engine def get_supabase_client() -> Client: global _supabase_client diff --git a/backend/modules/brain/entity/brain_entity.py b/backend/modules/brain/entity/brain_entity.py index 996a4abe9..a7ce90aa2 100644 --- a/backend/modules/brain/entity/brain_entity.py +++ b/backend/modules/brain/entity/brain_entity.py @@ -1,3 +1,4 @@ +from datetime import datetime from enum import Enum from typing import List, Optional from uuid import UUID @@ -26,7 +27,7 @@ class BrainEntity(BaseModel): max_tokens: Optional[int] = None status: Optional[str] = None prompt_id: Optional[UUID] = None - last_update: str + last_update: datetime brain_type: BrainType brain_definition: Optional[ApiBrainDefinitionEntity] = None connected_brains_ids: Optional[List[UUID]] = None diff --git a/backend/modules/brain/repository/brains.py b/backend/modules/brain/repository/brains.py index e8c7ac987..54b48c204 100644 --- a/backend/modules/brain/repository/brains.py +++ b/backend/modules/brain/repository/brains.py @@ -1,7 +1,9 @@ from uuid import UUID +from sqlalchemy import text + from logger import get_logger -from models.settings import get_embeddings, get_supabase_client +from models.settings import get_embeddings, get_pg_database_engine, get_supabase_client from modules.brain.dto.inputs import BrainUpdatableProperties from modules.brain.entity.brain_entity import BrainEntity, PublicBrain from modules.brain.repository.interfaces.brains_interface import BrainsInterface @@ -13,6 +15,8 @@ class Brains(BrainsInterface): def __init__(self): supabase_client = get_supabase_client() self.db = supabase_client + pg_engine = get_pg_database_engine() + self.pg_engine = pg_engine def create_brain(self, brain): embeddings = get_embeddings() @@ -54,27 +58,32 @@ class Brains(BrainsInterface): def update_brain_last_update_time(self, brain_id): try: - self.db.table("brains").update({"last_update": "now()"}).match( - {"brain_id": brain_id} - ).execute() + with self.pg_engine.begin() as connection: + query = """ + UPDATE brains + SET last_update = now() + WHERE brain_id = '{brain_id}' + """ + connection.execute(text(query.format(brain_id=brain_id))) except Exception as e: logger.error(e) def get_brain_details(self, brain_id): - response = ( - self.db.table("brains") - .select("*") - .filter("brain_id", "eq", str(brain_id)) - .execute() - ) - if response.data == []: + with self.pg_engine.begin() as connection: + query = """ + SELECT * FROM brains + WHERE brain_id = '{brain_id}' + """ + response = connection.execute(text(query.format(brain_id=brain_id))).fetchall() + if len(response) == 0: return None - return BrainEntity(**response.data[0]) + return BrainEntity(**response[0]._mapping) def delete_brain(self, brain_id: str): - results = ( - self.db.table("brains").delete().match({"brain_id": brain_id}).execute() - ) + with self.pg_engine.begin() as connection: + results = connection.execute( + text(f"DELETE FROM brains WHERE brain_id = '{brain_id}'") + ) return results @@ -100,14 +109,11 @@ class Brains(BrainsInterface): def get_brain_by_id(self, brain_id: UUID) -> BrainEntity | None: # TODO: merge this method with get_brain_details - response = ( - self.db.from_("brains") - .select("id:brain_id, name, *") - .filter("brain_id", "eq", brain_id) - .execute() - ).data + with self.pg_engine.begin() as connection: + response = connection.execute( + text(f"SELECT * FROM brains WHERE brain_id = '{brain_id}'") + ).fetchall() if len(response) == 0: return None - - return BrainEntity(**response[0]) + return BrainEntity(**response[0]._mapping) \ No newline at end of file