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):
This commit is contained in:
Damien Mourot 2024-04-28 22:27:50 -10:00 committed by GitHub
parent 89f86d4e68
commit fb3aae27f1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 41 additions and 24 deletions

View File

@ -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

View File

@ -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

View File

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