mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-14 17:03:29 +03:00
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:
parent
89f86d4e68
commit
fb3aae27f1
@ -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
|
||||
|
@ -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
|
||||
|
@ -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)
|
Loading…
Reference in New Issue
Block a user