Skip to content

Commit

Permalink
feat: Introduce repository pattern to prepare adding other database p…
Browse files Browse the repository at this point in the history 
…roviders (#646)

* add sqlalchemy models

* add neon settings

* add insert brain

* abstract supabase from Brain class

* abstract supabase from Brain class

* abstract supabase from /models

* update Database to Repository

* update neon_tables to pg_tables

* update chat, api-key and message

* update vector class

* update settings

* update env vars for test

* Update backend-tests.yml

* fix test

* fix fetch_user_requests_count()

* fix fetch_user_requests_count()

* fix increment_user_request_count

* fix increment_user_request_count

* fix asset upload_response message

* fix pyright

* fix brain_subscription

* fix brain_subscription

* fix brain_subscription

* fix get user request stat

* update create_brain_user

* add delete brain vector and user

* add delete brain vector and user

* correctly call function

---------

Co-authored-by: Noé Pion <noe.pion@onfido.com>
Co-authored-by: raoufchebri <raouf@chebri.com>
Co-authored-by: Stan Girard <girard.stanislas@gmail.com>
  • Loading branch information
@Noezor @raoufchebri @StanGirard
4 people authored Aug 1, 2023
1 parent d6ebb00 commit 303ba72
Show file tree
Hide file tree
Showing 33 changed files with 1,189 additions and 422 deletions.
1 change: 1 addition & 0 deletions .backend_env.example
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
SUPABASE_URL=<change-me>
SUPABASE_SERVICE_KEY=<change-me>
PG_DATABASE_URL=<change-me>
OPENAI_API_KEY=<change-me>
ANTHROPIC_API_KEY=null
JWT_SECRET_KEY=<change-me>
Expand Down
35 changes: 28 additions & 7 deletions .github/workflows/backend-tests.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,33 @@ name: Backend Tests

on:
push:
branches: [ main ]
branches: [main]
paths:
- 'backend/core/**'
pull_request:
branches: [ main ]
branches: [main]
paths:
- 'backend/core/**'
- 'backend/core/**'
workflow_dispatch:
inputs:
logLevel:
description: 'Log level'
required: true
default: 'warning'
type: choice
options:
- info
- warning
- debug
tags:
description: 'Test scenario tags'
required: false
type: boolean
environment:
description: 'Environment to run tests against'
type: environment
required: false

jobs:
build:
runs-on: ubuntu-latest
Expand All @@ -18,15 +38,15 @@ jobs:
environment: preview
strategy:
matrix:
python-version: [ "3.11"]
python-version: ["3.11"]

steps:
- uses: actions/checkout@v3
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
cache: 'pip'
cache: "pip"
- name: Install dependencies
run: |
python -m pip install --upgrade pip
Expand All @@ -37,13 +57,14 @@ jobs:
env:
SUPABASE_URL: ${{secrets.SUPABASE_URL}}
SUPABASE_SERVICE_KEY: ${{secrets.SUPABASE_SERVICE_KEY}}
PG_DATABASE_URL: ${{secrets.PG_DATABASE_URL}}
OPENAI_API_KEY: ${{secrets.OPENAI_API_KEY}}
ANTHROPIC_API_KEY: ${{secrets.ANTHROPIC_API_KEY}}
JWT_SECRET_KEY: ${{secrets.JWT_SECRET_KEY}}
CI_TEST_API_KEY: ${{secrets.CI_TEST_API_KEY}}
CI_TEST_API_KEY: ${{secrets.CI_TEST_API_KEY}}
run: |
python -m pytest tests/
- name: Static type checking with pyright
run: |
pyright
27 changes: 4 additions & 23 deletions backend/core/auth/api_key_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,8 @@ async def verify_api_key(
# Use UTC time to avoid timezone issues
current_date = datetime.utcnow().date()
commons = common_dependencies()
result = (
commons["supabase"]
.table("api_keys")
.select("api_key", "creation_time")
.filter("api_key", "eq", api_key)
.filter("is_active", "eq", True)
.execute()
)
result = commons["db"].get_active_api_key(api_key)

if result.data is not None and len(result.data) > 0:
api_key_creation_date = datetime.strptime(
result.data[0]["creation_time"], "%Y-%m-%dT%H:%M:%S"
Expand All @@ -42,28 +36,15 @@ async def get_user_from_api_key(
commons = common_dependencies()

# Lookup the user_id from the api_keys table
user_id_data = (
commons["supabase"]
.table("api_keys")
.select("user_id")
.filter("api_key", "eq", api_key)
.execute()
)
user_id_data = commons["db"].get_user_id_by_api_key(api_key)

if not user_id_data.data:
raise HTTPException(status_code=400, detail="Invalid API key.")

user_id = user_id_data.data[0]["user_id"]

# Lookup the email from the users table. Todo: remove and use user_id for credentials
user_email_data = (
commons["supabase"]
.table("users")
.select("email")
.filter("user_id", "eq", user_id)
.execute()
)

user_email_data = commons["db"].get_user_email(user_id)
email = user_email_data.data[0]["email"] if user_email_data.data else None

return User(email=email, id=user_id)
209 changes: 25 additions & 184 deletions backend/core/models/brains.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -78,31 +78,10 @@ def delete_user_from_brain(self, user_id):
).execute()

def get_user_brains(self, user_id):
response = (
self.commons["supabase"]
.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)")
.filter("user_id", "eq", user_id)
.execute()
)
user_brains = []
for item in response.data:
user_brains.append(item["brains"])
user_brains[-1]["rights"] = item["rights"]
return user_brains
return self.commons["db"].get_user_brains(user_id)

def get_brain_for_user(self, user_id):
response = (
self.commons["supabase"]
.from_("brains_users")
.select("id:brain_id, rights, brains (id: brain_id, name)")
.filter("user_id", "eq", user_id)
.filter("brain_id", "eq", self.id)
.execute()
)
if len(response.data) == 0:
return None
return response.data[0]
return self.commons["db"].get_brain_for_user(user_id, self.id)

def get_brain_details(self):
response = (
Expand All @@ -117,78 +96,23 @@ def get_brain_details(self):
return response.data[0]

def delete_brain(self, user_id):
results = (
self.commons["supabase"]
.table("brains_users")
.select("*")
.match({"brain_id": self.id, "user_id": user_id, "rights": "Owner"})
.execute()
)
results = self.commons["db"].delete_brain_user_by_id(user_id, self.id)

if len(results.data) == 0:
return {"message": "You are not the owner of this brain."}
else:
results = (
self.commons["supabase"]
.table("brains_vectors")
.delete()
.match({"brain_id": self.id})
.execute()
)

results = (
self.commons["supabase"]
.table("brains_users")
.delete()
.match({"brain_id": self.id})
.execute()
)

results = (
self.commons["supabase"]
.table("brains")
.delete()
.match({"brain_id": self.id})
.execute()
)
self.commons["db"].delete_brain_vector(self.id)
self.commons["db"].delete_brain_user(self.id)
self.commons["db"].delete_brain(self.id)

def create_brain(self):
commons = common_dependencies()
response = (
commons["supabase"]
.table("brains")
.insert(
{
"name": self.name,
"description": self.description,
"temperature": self.temperature,
"model": self.model,
"max_tokens": self.max_tokens,
"openai_api_key": self.openai_api_key,
"status": self.status,
}
)
.execute()
)

response = self.commons["db"].create_brain(self.name)
self.id = response.data[0]["brain_id"]
return response.data

def create_brain_user(self, user_id: UUID, rights, default_brain: bool):
commons = common_dependencies()
response = (
commons["supabase"]
.table("brains_users")
.insert(
{
"brain_id": str(self.id),
"user_id": str(user_id),
"rights": rights,
"default_brain": default_brain,
}
)
.execute()
)

def create_brain_user(self, user_id: UUID, rights, default_brain):
response = self.commons["db"].create_brain_user(user_id=user_id, brain_id=self.id, rights=rights, default_brain=default_brain)
self.id = response.data[0]["brain_id"]
return response.data

def set_as_default_brain_for_user(self, user: User):
Expand All @@ -204,128 +128,45 @@ def set_as_default_brain_for_user(self, user: User):
).match({"brain_id": self.id, "user_id": user.id}).execute()

def create_brain_vector(self, vector_id, file_sha1):
response = (
self.commons["supabase"]
.table("brains_vectors")
.insert(
{
"brain_id": str(self.id),
"vector_id": str(vector_id),
"file_sha1": file_sha1,
}
)
.execute()
)
return response.data
return self.commons["db"].create_brain_vector(self.id, vector_id, file_sha1)

def get_vector_ids_from_file_sha1(self, file_sha1: str):
# move to vectors class
vectorsResponse = (
self.commons["supabase"]
.table("vectors")
.select("id")
.filter("metadata->>file_sha1", "eq", file_sha1)
.execute()
)
return vectorsResponse.data
return self.commons["db"].get_vector_ids_from_file_sha1(file_sha1)

def update_brain_fields(self):
self.commons["supabase"].table("brains").update(
{
"name": self.name,
"description": self.description,
"model": self.model,
"temperature": self.temperature,
"max_tokens": self.max_tokens,
"openai_api_key": self.openai_api_key,
"status": self.status,
}
).match({"brain_id": self.id}).execute()
return self.commons["db"].update_brain_fields(brain_id=self.id, brain_name=self.name)

def update_brain_with_file(self, file_sha1: str):
# not used
vector_ids = self.get_vector_ids_from_file_sha1(file_sha1)
for vector_id in vector_ids:
self.create_brain_vector(vector_id, file_sha1)

def get_unique_brain_files(self):
"""
Retrieve unique brain data (i.e. uploaded files and crawled websites).
"""

response = (
self.commons["supabase"]
.from_("brains_vectors")
.select("vector_id")
.filter("brain_id", "eq", self.id)
.execute()
)

vector_ids = [item["vector_id"] for item in response.data]

if len(vector_ids) == 0:
return []

vector_ids = self.commons["db"].get_brain_vector_ids(self.id)
self.files = get_unique_files_from_vector_ids(vector_ids)

return self.files

def delete_file_from_brain(self, file_name: str):
# First, get the vector_ids associated with the file_name
vector_response = (
self.commons["supabase"]
.table("vectors")
.select("id")
.filter("metadata->>file_name", "eq", file_name)
.execute()
)
vector_ids = [item["id"] for item in vector_response.data]

# For each vector_id, delete the corresponding entry from the 'brains_vectors' table
for vector_id in vector_ids:
self.commons["supabase"].table("brains_vectors").delete().filter(
"vector_id", "eq", vector_id
).filter("brain_id", "eq", self.id).execute()

# Check if the vector is still associated with any other brains
associated_brains_response = (
self.commons["supabase"]
.table("brains_vectors")
.select("brain_id")
.filter("vector_id", "eq", vector_id)
.execute()
)
associated_brains = [
item["brain_id"] for item in associated_brains_response.data
]

# If the vector is not associated with any other brains, delete it from 'vectors' table
if not associated_brains:
self.commons["supabase"].table("vectors").delete().filter(
"id", "eq", vector_id
).execute()

return {"message": f"File {file_name} in brain {self.id} has been deleted."}
return self.commons["db"].delete_file_from_brain(self.id, file_name)


def get_default_user_brain(user: User):
commons = common_dependencies()
response = (
commons["supabase"]
.from_("brains_users")
.select("brain_id")
.filter("user_id", "eq", user.id)
.filter("default_brain", "eq", True)
.execute()
)
response = commons["db"].get_default_user_brain_id(user.id)

logger.info("Default brain response:", response.data)
default_brain_id = response.data[0]["brain_id"] if response.data else None

logger.info(f"Default brain id: {default_brain_id}")

if default_brain_id:
brain_response = (
commons["supabase"]
.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", default_brain_id)
.execute()
)

brain_response = commons["db"].get_brain_by_id(default_brain_id)
return brain_response.data[0] if brain_response.data else None


Expand Down
Loading

0 comments on commit 303ba72

Please sign in to comment.