QuivrHQ/quivr
1
1
Fork 0
mirror of https://github.com/QuivrHQ/quivr.git synced 2024-09-11 14:36:35 +03:00
Code Issues Packages Projects Releases Wiki Activity

Feat/multiple brains files (#361)

Browse Source
This commit is contained in:
Zineb El Bachiri 2023-06-28 19:39:27 +02:00 committed by GitHub
parent e79da8e3cd
commit ccdc5bb7a7
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
51 changed files with 1205 additions and 668 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
backend/auth/auth_bearer.py
View File

@ -55,4 +55,4 @@ class AuthBearer(HTTPBearer):
def get_current_user(credentials: dict = Depends(AuthBearer())) -> User:
return User(email=credentials.get("email", "none"))
return User(email=credentials.get("email", "none"), id=credentials.get("sub", "none"))

11
backend/llm/BrainPickingOpenAIFunctions/BrainPickingOpenAIFunctions.py
View File

@ -17,7 +17,7 @@ class BrainPickingOpenAIFunctions(BrainPicking):
DEFAULT_MAX_TOKENS = 256
openai_client: ChatOpenAI = None
user_email: str = None
brain_id: str = None
def __init__(
self,
@ -25,20 +25,19 @@ class BrainPickingOpenAIFunctions(BrainPicking):
chat_id: str,
temperature: float,
max_tokens: int,
user_email: str,
brain_id: str,
user_openai_api_key: str,
) -> None:
# Call the constructor of the parent class (BrainPicking)
super().__init__(
model=model,
user_id=user_email,
chat_id=chat_id,
max_tokens=max_tokens,
user_openai_api_key=user_openai_api_key,
temperature=temperature,
brain_id=str(brain_id),
)
self.openai_client = ChatOpenAI(openai_api_key=self.settings.openai_api_key)
self.user_email = user_email
def _get_model_response(
self,
@ -86,10 +85,10 @@ class BrainPickingOpenAIFunctions(BrainPicking):
self.supabase_client,
self.embeddings,
table_name="vectors",
user_id=self.user_email,
brain_id=self.brain_id,
)
return vector_store.similarity_search(query=question, user_id=self.user_email)
return vector_store.similarity_search(query=question)
def _construct_prompt(
self, question: str, useContext: bool = False, useHistory: bool = False

22
backend/llm/PrivateBrainPicking.py
View File

@ -1,25 +1,9 @@
from typing import Any, Dict
# Importing various modules and classes from a custom library 'langchain' likely used for natural language processing
from langchain.chains import ConversationalRetrievalChain, LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import GPT4All
from langchain.llms.base import LLM
from langchain.memory import ConversationBufferMemory
from llm.brainpicking import BrainPicking
from llm.prompt.CONDENSE_PROMPT import CONDENSE_QUESTION_PROMPT
from logger import get_logger
from models.settings import BrainSettings # Importing settings related to the 'brain'
from models.settings import LLMSettings # For type hinting
from pydantic import BaseModel # For data validation and settings management
from repository.chat.get_chat_history import get_chat_history
from supabase import Client # For interacting with Supabase database
from supabase import create_client
from vectorstore.supabase import (
CustomSupabaseVectorStore,
) # Custom class for handling vector storage with Supabase
logger = get_logger(__name__)
@ -32,8 +16,8 @@ class PrivateBrainPicking(BrainPicking):
def __init__(
self,
model: str,
user_id: str,
chat_id: str,
brain_id:str,
temperature: float,
max_tokens: int,
user_openai_api_key: str,
@ -41,13 +25,13 @@ class PrivateBrainPicking(BrainPicking):
"""
Initialize the PrivateBrainPicking class by calling the parent class's initializer.
:param model: Language model name to be used.
:param user_id: The user id to be used for CustomSupabaseVectorStore.
:param brain_id: The user id to be used for CustomSupabaseVectorStore.
:return: PrivateBrainPicking instance
"""
# Call the parent class's initializer
super().__init__(
model=model,
user_id=user_id,
brain_id=brain_id,
chat_id=chat_id,
max_tokens=max_tokens,
temperature=temperature,

18
backend/llm/brainpicking.py
View File

@ -5,20 +5,20 @@ from langchain.chains import ConversationalRetrievalChain, LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import GPT4All
from langchain.llms.base import LLM
from langchain.memory import ConversationBufferMemory
from llm.prompt.CONDENSE_PROMPT import CONDENSE_QUESTION_PROMPT
from logger import get_logger
from models.settings import BrainSettings # Importing settings related to the 'brain'
from models.settings import \
BrainSettings # Importing settings related to the 'brain'
from models.settings import LLMSettings # For type hinting
from pydantic import BaseModel # For data validation and settings management
from repository.chat.get_chat_history import get_chat_history
from vectorstore.supabase import \
CustomSupabaseVectorStore # Custom class for handling vector storage with Supabase
from supabase import Client # For interacting with Supabase database
from supabase import create_client
from vectorstore.supabase import (
CustomSupabaseVectorStore,
) # Custom class for handling vector storage with Supabase
logger = get_logger(__name__)
@ -76,7 +76,7 @@ class BrainPicking(BaseModel):
def __init__(
self,
model: str,
user_id: str,
brain_id: str,
temperature: float,
chat_id: str,
max_tokens: int,
@ -85,12 +85,12 @@ class BrainPicking(BaseModel):
"""
Initialize the BrainPicking class by setting embeddings, supabase client, vector store, language model and chains.
:param model: Language model name to be used.
:param user_id: The user id to be used for CustomSupabaseVectorStore.
:param user_brain_idid: The brain id to be used for CustomSupabaseVectorStore.
:return: BrainPicking instance
"""
super().__init__(
model=model,
user_id=user_id,
brain_id=brain_id,
chat_id=chat_id,
max_tokens=max_tokens,
temperature=temperature,
@ -110,7 +110,7 @@ class BrainPicking(BaseModel):
self.supabase_client,
self.embeddings,
table_name="vectors",
user_id=user_id,
brain_id=brain_id,
)
self.llm = self._determine_llm(

1
backend/main.py
View File

@ -27,7 +27,6 @@ max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200)
async def startup_event():
pypandoc.download_pandoc()
app.include_router(brain_router)
app.include_router(chat_router)
app.include_router(crawl_router)

157
backend/models/brains.py
View File

@ -1,18 +1,22 @@
from typing import Optional
import os
from typing import Any, List, Optional
from uuid import UUID
from models.settings import CommonsDep, common_dependencies
from models.users import User
from pydantic import BaseModel
class Brain(BaseModel):
brain_id: Optional[UUID]
name: Optional[str] = "New Brain"
status: Optional[str] = "public"
id: Optional[UUID] = None
name: Optional[str] = "Default brain"
status: Optional[str]= "public"
model: Optional[str] = "gpt-3.5-turbo-0613"
temperature: Optional[float] = 0.0
max_tokens: Optional[int] = 256
file_sha1: Optional[str] = ""
brain_size: Optional[float] = 0.0
max_brain_size: Optional[int] = int(os.getenv("MAX_BRAIN_SIZE", 0))
files: List[Any] = []
_commons: Optional[CommonsDep] = None
class Config:
@ -24,6 +28,19 @@ class Brain(BaseModel):
self.__class__._commons = common_dependencies()
return self._commons
@property
def brain_size(self):
self.get_unique_brain_files()
current_brain_size = sum(float(doc['size']) for doc in self.files)
print('current_brain_size', current_brain_size)
return current_brain_size
@property
def remaining_brain_size(self):
return float(self.max_brain_size) - self.brain_size
@classmethod
def create(cls, *args, **kwargs):
commons = common_dependencies()
@ -39,56 +56,47 @@ class Brain(BaseModel):
)
return [item["brains"] for item in response.data]
def get_brain(self):
response = (
self.commons["supabase"]
.from_("brains")
.select("brainId:brain_id, brainName:brain_name")
.filter("brain_id", "eq", self.brain_id)
.execute()
)
return response.data
def get_brain_details(self):
response = (
self.commons["supabase"]
.from_("brains")
.select("id:brain_id, name, *")
.filter("brain_id", "eq", self.brain_id)
.filter("brain_id", "eq", self.id)
.execute()
)
return response.data
def delete_brain(self):
self.commons["supabase"].table("brains").delete().match(
{"brain_id": self.brain_id}
{"brain_id": self.id}
).execute()
@classmethod
def create_brain(cls, name):
def create_brain(self):
commons = common_dependencies()
response = commons["supabase"].table("brains").insert({"name": name}).execute()
response = commons["supabase"].table("brains").insert({"name": self.name}).execute()
# set the brainId with response.data
self.id = response.data[0]['brain_id']
return response.data
def create_brain_user(self, brain_id, user_id, rights):
response = (
self.commons["supabase"]
.table("brains_users")
.insert({"brain_id": brain_id, "user_id": user_id, "rights": rights})
.execute()
)
def create_brain_user(self, user_id : UUID, rights, default_brain):
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()
return response.data
def create_brain_vector(self, vector_id):
response = (
self.commons["supabase"]
.table("brains_users")
.insert({"brain_id": self.brain_id, "vector_id": vector_id})
.table("brains_vectors")
.insert({"brain_id": str(self.id), "vector_id": str(vector_id)})
.execute()
)
return response.data
def get_vector_ids_from_file_sha1(self, file_sha1: str):
# move to vectors class
vectorsResponse = (
self.commons["supabase"]
.table("vectors")
@ -100,10 +108,99 @@ class Brain(BaseModel):
def update_brain_fields(self):
self.commons["supabase"].table("brains").update({"name": self.name}).match(
{"brain_id": self.brain_id}
{"brain_id": self.id}
).execute()
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)
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]
print('vector_ids', vector_ids)
if len(vector_ids) == 0:
return []
self.files = self.get_unique_files_from_vector_ids(vector_ids)
print('unique_files', self.files)
return self.files
def get_unique_files_from_vector_ids(self, vectors_ids : List[int]):
# Move into Vectors class
"""
Retrieve unique user data vectors.
"""
vectors_response = self.commons['supabase'].table("vectors").select(
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
.filter("id", "in", tuple(vectors_ids))\
.execute()
documents = vectors_response.data # Access the data from the response
# Convert each dictionary to a tuple of items, then to a set to remove duplicates, and then back to a dictionary
unique_files = [dict(t) for t in set(tuple(d.items()) for d in documents)]
return unique_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."}
def get_default_user_brain(user: User):
commons = common_dependencies()
response = (
commons["supabase"]
.from_("brains_users") # I'm assuming this is the correct table
.select("brain_id")
.filter("user_id", "eq", user.id)
.filter("default_brain", "eq", True) # Assuming 'default' is the correct column name
.execute()
)
default_brain_id = response.data[0]["brain_id"] if response.data else None
print(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()
)
return brain_response.data[0] if brain_response.data else None
return None

109
backend/models/files.py Normal file
View File

@ -0,0 +1,109 @@
import os
import tempfile
from typing import Any, Optional
from uuid import UUID
from fastapi import UploadFile
from langchain.text_splitter import RecursiveCharacterTextSplitter
from logger import get_logger
from models.settings import CommonsDep, common_dependencies
from pydantic import BaseModel
from utils.file import compute_sha1_from_file
logger = get_logger(__name__)
class File(BaseModel):
id: Optional[UUID] = None
file: Optional[UploadFile]
file_name: Optional[str] = ""
file_size: Optional[int] = ""
file_sha1: Optional[str] = ""
vectors_ids: Optional[int]=[]
file_extension: Optional[str] = ""
content: Optional[Any]= None
chunk_size: int = 500
chunk_overlap: int= 0
documents: Optional[Any]= None
_commons: Optional[CommonsDep] = None
def __init__(self, **kwargs):
super().__init__(**kwargs)
if self.file:
self.file_name = self.file.filename
self.file_size = self.file.file._file.tell()
self.file_extension = os.path.splitext(self.file.filename)[-1].lower()
async def compute_file_sha1(self):
with tempfile.NamedTemporaryFile(delete=False, suffix=self.file.filename) as tmp_file:
await self.file.seek(0)
self.content = await self.file.read()
tmp_file.write(self.content)
tmp_file.flush()
self.file_sha1 = compute_sha1_from_file(tmp_file.name)
os.remove(tmp_file.name)
def compute_documents(self, loader_class):
logger.info(f"Computing documents from file {self.file_name}")
documents = []
with tempfile.NamedTemporaryFile(delete=False, suffix=self.file.filename) as tmp_file:
tmp_file.write(self.content)
tmp_file.flush()
loader = loader_class(tmp_file.name)
documents = loader.load()
print("documents", documents)
os.remove(tmp_file.name)
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
)
self.documents = text_splitter.split_documents(documents)
print(self.documents)
def set_file_vectors_ids(self):
commons = common_dependencies()
response = (
commons["supabase"].table("vectors")
.select("id")
.filter("metadata->>file_sha1", "eq", self.file_sha1)
.execute()
)
self.vectors_ids = response.data
return
def file_already_exists(self, brain_id):
commons = common_dependencies()
self.set_file_vectors_ids()
print("file_sha1", self.file_sha1)
print("vectors_ids", self.vectors_ids)
print("len(vectors_ids)", len(self.vectors_ids))
if len(self.vectors_ids) == 0:
return False
for vector in self.vectors_ids:
response = (
commons["supabase"].table("brains_vectors")
.select("brain_id, vector_id")
.filter("brain_id", "eq", brain_id)
.filter("vector_id", "eq", vector['id'])
.execute()
)
print("response.data", response.data)
if len(response.data) == 0:
return False
return True
def file_is_empty(self):
return self.file.file._file.tell() < 1

51
backend/models/users.py
View File

@ -1,5 +1,54 @@
from uuid import UUID
from logger import get_logger
from models.settings import common_dependencies
from pydantic import BaseModel
logger = get_logger(__name__)
class User(BaseModel):
id: UUID
email: str
user_openai_api_key: str = None
requests_count: int = 0
user_openai_api_key: str = None
# [TODO] Rename the user table and its references to 'user_usage'
def create_user( self,date):
commons = common_dependencies()
logger.info(f"New user entry in db document for user {self.email}")
return(commons['supabase'].table("users").insert(
{"user_id": self.id, "email": self.email, "date": date, "requests_count": 1}).execute())
def get_user_request_stats(self):
commons = common_dependencies()
requests_stats = commons['supabase'].from_('users').select(
'*').filter("user_id", "eq", self.id).execute()
return requests_stats.data
def fetch_user_requests_count(self, date):
commons = common_dependencies()
response = (
commons["supabase"]
.from_("users")
.select("*")
.filter("user_id", "eq", self.id)
.filter("date", "eq", date)
.execute()
)
userItem = next(iter(response.data or []), {"requests_count": 0})
return userItem["requests_count"]
def increment_user_request_count(self, date):
commons = common_dependencies()
requests_count = self.fetch_user_requests_count(date) + 1
logger.info(f"User {self.email} request count updated to {requests_count}")
commons['supabase'].table("users").update(
{ "requests_count": requests_count}).match({"user_id": self.id, "date": date}).execute()
self.requests_count = requests_count

37
backend/parsers/audio.py
View File

@ -1,51 +1,29 @@
import os
import tempfile
import time
from io import BytesIO
from tempfile import NamedTemporaryFile
import openai
from fastapi import UploadFile
from langchain.document_loaders import TextLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from models.files import File
from models.settings import CommonsDep
from utils.file import compute_sha1_from_content
# # Create a function to transcribe audio using Whisper
# def _transcribe_audio(api_key, audio_file, stats_db):
# openai.api_key = api_key
# transcript = ""
# with BytesIO(audio_file.read()) as audio_bytes:
# # Get the extension of the uploaded file
# file_extension = os.path.splitext(audio_file.name)[-1]
# # Create a temporary file with the uploaded audio data and the correct extension
# with tempfile.NamedTemporaryFile(delete=True, suffix=file_extension) as temp_audio_file:
# temp_audio_file.write(audio_bytes.read())
# temp_audio_file.seek(0) # Move the file pointer to the beginning of the file
# transcript = openai.Audio.translate("whisper-1", temp_audio_file)
# return transcript
# async def process_audio(upload_file: UploadFile, stats_db):
async def process_audio(commons: CommonsDep, upload_file: UploadFile, enable_summarization: bool, user, user_openai_api_key):
async def process_audio(commons: CommonsDep, file: File, enable_summarization: bool, user, user_openai_api_key):
temp_filename = None
file_sha = ""
dateshort = time.strftime("%Y%m%d-%H%M%S")
file_meta_name = f"audiotranscript_{dateshort}.txt"
# uploaded file to file object
# use this for whisper
openai_api_key = os.environ.get("OPENAI_API_KEY")
if user_openai_api_key:
openai_api_key = user_openai_api_key
try:
# Here, we're writing the uploaded file to a temporary file, so we can use it with your existing code.
upload_file = file.file
with tempfile.NamedTemporaryFile(delete=False, suffix=upload_file.filename) as tmp_file:
await upload_file.seek(0)
content = await upload_file.read()
@ -61,7 +39,6 @@ async def process_audio(commons: CommonsDep, upload_file: UploadFile, enable_sum
file_sha = compute_sha1_from_content(transcript.text.encode("utf-8"))
file_size = len(transcript.text.encode("utf-8"))
# Load chunk size and overlap from sidebar
chunk_size = 500
chunk_overlap = 0
@ -72,12 +49,8 @@ async def process_audio(commons: CommonsDep, upload_file: UploadFile, enable_sum
docs_with_metadata = [Document(page_content=text, metadata={"file_sha1": file_sha, "file_size": file_size, "file_name": file_meta_name,
"chunk_size": chunk_size, "chunk_overlap": chunk_overlap, "date": dateshort}) for text in texts]
# if st.secrets.self_hosted == "false":
# add_usage(stats_db, "embedding", "audio", metadata={"file_name": file_meta_name,"file_type": ".txt", "chunk_size": chunk_size, "chunk_overlap": chunk_overlap})
commons.documents_vector_store.add_documents(docs_with_metadata)
finally:
if temp_filename and os.path.exists(temp_filename):
os.remove(temp_filename)
return documents_vector_store
os.remove(temp_filename)

92
backend/parsers/common.py
View File

@ -1,100 +1,44 @@
# from stats import add_usage
import asyncio
import os
import tempfile
import time
from typing import Optional
from fastapi import UploadFile
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from models.brains import Brain
from models.files import File
from models.settings import CommonsDep
from utils.file import compute_sha1_from_content, compute_sha1_from_file
from utils.vectors import Neurons, create_summary
from utils.vectors import Neurons
async def process_file(
commons: CommonsDep,
file: UploadFile,
file: File,
loader_class,
file_suffix,
enable_summarization,
user,
brain_id,
user_openai_api_key,
):
documents = []
file_name = file.filename
file_size = file.file._file.tell() # Getting the size of the file
dateshort = time.strftime("%Y%m%d")
# Here, we're writing the uploaded file to a temporary file, so we can use it with your existing code.
with tempfile.NamedTemporaryFile(delete=False, suffix=file.filename) as tmp_file:
await file.seek(0)
content = await file.read()
tmp_file.write(content)
tmp_file.flush()
file.compute_documents(loader_class)
loader = loader_class(tmp_file.name)
documents = loader.load()
# Ensure this function works with FastAPI
file_sha1 = compute_sha1_from_file(tmp_file.name)
os.remove(tmp_file.name)
chunk_size = 500
chunk_overlap = 0
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
chunk_size=chunk_size, chunk_overlap=chunk_overlap
)
documents = text_splitter.split_documents(documents)
for doc in documents:
for doc in file.documents:
metadata = {
"file_sha1": file_sha1,
"file_size": file_size,
"file_name": file_name,
"chunk_size": chunk_size,
"chunk_overlap": chunk_overlap,
"file_sha1": file.file_sha1,
"file_size": file.file_size,
"file_name": file.file_name,
"chunk_size": file.chunk_size,
"chunk_overlap": file.chunk_overlap,
"date": dateshort,
"summarization": "true" if enable_summarization else "false",
}
doc_with_metadata = Document(
page_content=doc.page_content, metadata=metadata)
neurons = Neurons(commons=commons)
neurons.create_vector(user.email, doc_with_metadata, user_openai_api_key)
created_vector = neurons.create_vector(doc_with_metadata, user_openai_api_key)
# add_usage(stats_db, "embedding", "audio", metadata={"file_name": file_meta_name,"file_type": ".txt", "chunk_size": chunk_size, "chunk_overlap": chunk_overlap})
# Remove the enable_summarization and ids
if enable_summarization and ids and len(ids) > 0:
create_summary(
commons, document_id=ids[0], content=doc.page_content, metadata=metadata
)
created_vector_id = created_vector[0]
brain = Brain(id=brain_id)
brain.create_brain_vector(created_vector_id)
return
async def file_already_exists(supabase, file, user):
# TODO: user brain id instead of user
file_content = await file.read()
file_sha1 = compute_sha1_from_content(file_content)
response = (
supabase.table("vectors")
.select("id")
.filter("metadata->>file_sha1", "eq", file_sha1)
.filter("user_id", "eq", user.email)
.execute()
)
return len(response.data) > 0
async def file_already_exists_from_content(supabase, file_content, user):
# TODO: user brain id instead of user
file_sha1 = compute_sha1_from_content(file_content)
response = (
supabase.table("vectors")
.select("id")
.filter("metadata->>file_sha1", "eq", file_sha1)
.filter("user_id", "eq", user.email)
.execute()
)
return len(response.data) > 0

9
backend/parsers/csv.py
View File

@ -1,5 +1,5 @@
from fastapi import UploadFile
from langchain.document_loaders.csv_loader import CSVLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
@ -7,17 +7,16 @@ from .common import process_file
def process_csv(
commons: CommonsDep,
file: UploadFile,
file: File,
enable_summarization,
user,
brain_id,
user_openai_api_key,
):
return process_file(
commons,
file,
CSVLoader,
".csv",
enable_summarization,
user,
brain_id,
user_openai_api_key,
)

6
backend/parsers/docx.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import Docx2txtLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_docx(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, Docx2txtLoader, ".docx", enable_summarization, user, user_openai_api_key)
def process_docx(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, Docx2txtLoader, enable_summarization, brain_id, user_openai_api_key)

6
backend/parsers/epub.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders.epub import UnstructuredEPubLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_epub(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, UnstructuredEPubLoader, ".epub", enable_summarization, user, user_openai_api_key)
def process_epub(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, UnstructuredEPubLoader, enable_summarization, brain_id, user_openai_api_key)

20
backend/parsers/github.py
View File

@ -4,13 +4,14 @@ import time
from langchain.document_loaders import GitLoader
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from models.brains import Brain
from models.files import File
from models.settings import CommonsDep
from parsers.common import file_already_exists_from_content
from utils.file import compute_sha1_from_content
from utils.vectors import Neurons
async def process_github(commons: CommonsDep, repo, enable_summarization, user, supabase, user_openai_api_key):
async def process_github(commons: CommonsDep, repo, enable_summarization, brain_id, user_openai_api_key):
random_dir_name = os.urandom(16).hex()
dateshort = time.strftime("%Y%m%d")
loader = GitLoader(
@ -42,11 +43,22 @@ async def process_github(commons: CommonsDep, repo, enable_summarization, user,
}
doc_with_metadata = Document(
page_content=doc.page_content, metadata=metadata)
exist = await file_already_exists_from_content(supabase, doc.page_content.encode("utf-8"), user)
file = File(file_sha1 = compute_sha1_from_content(doc.page_content.encode("utf-8")))
exist = file.file_already_exists(brain_id)
if not exist:
neurons = Neurons(commons=commons)
neurons.create_vector(user.email, doc_with_metadata, user_openai_api_key)
created_vector = neurons.create_vector(doc_with_metadata, user_openai_api_key)
created_vector_id = created_vector[0]
brain = Brain(id=brain_id)
brain.create_brain_vector(created_vector_id)
print("Created vector for ", doc.metadata["file_name"])
# add created_vector x brains in db
return {"message": f"✅ Github with {len(documents)} files has been uploaded.", "type": "success"}

8
backend/parsers/html.py
View File

@ -1,18 +1,16 @@
import os
import re
import tempfile
import unicodedata
import requests
from fastapi import UploadFile
from langchain.document_loaders import UnstructuredHTMLLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_html(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, UnstructuredHTMLLoader, ".html", enable_summarization, user, user_openai_api_key)
def process_html(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, UnstructuredHTMLLoader, enable_summarization, brain_id, user_openai_api_key)
def get_html(url):

6
backend/parsers/markdown.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import UnstructuredMarkdownLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_markdown(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, UnstructuredMarkdownLoader, ".md", enable_summarization, user, user_openai_api_key)
def process_markdown(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, UnstructuredMarkdownLoader, enable_summarization, brain_id, user_openai_api_key)

6
backend/parsers/notebook.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import NotebookLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_ipnyb(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, NotebookLoader, "ipynb", enable_summarization, user, user_openai_api_key)
def process_ipnyb(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, NotebookLoader, enable_summarization, brain_id, user_openai_api_key)

8
backend/parsers/odt.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import UnstructuredODTLoader
from langchain.document_loaders import PyMuPDFLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_odt(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, UnstructuredODTLoader, ".odt", enable_summarization, user, user_openai_api_key)
def process_odt(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, PyMuPDFLoader, enable_summarization, brain_id, user_openai_api_key)

6
backend/parsers/pdf.py
View File

@ -1,10 +1,10 @@
from fastapi import UploadFile
from langchain.document_loaders import PyMuPDFLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_pdf(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, PyMuPDFLoader, ".pdf", enable_summarization, user, user_openai_api_key)
def process_pdf(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, PyMuPDFLoader, enable_summarization, brain_id, user_openai_api_key)

6
backend/parsers/powerpoint.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import UnstructuredPowerPointLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
def process_powerpoint(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return process_file(commons, file, UnstructuredPowerPointLoader, ".pptx", enable_summarization, user, user_openai_api_key)
def process_powerpoint(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return process_file(commons, file, UnstructuredPowerPointLoader, enable_summarization, brain_id, user_openai_api_key)

6
backend/parsers/txt.py
View File

@ -1,9 +1,9 @@
from fastapi import UploadFile
from langchain.document_loaders import TextLoader
from models.files import File
from models.settings import CommonsDep
from .common import process_file
async def process_txt(commons: CommonsDep, file: UploadFile, enable_summarization, user, user_openai_api_key):
return await process_file(commons, file, TextLoader, ".txt", enable_summarization, user,user_openai_api_key)
async def process_txt(commons: CommonsDep, file: File, enable_summarization, brain_id, user_openai_api_key):
return await process_file(commons, file, TextLoader, enable_summarization, brain_id,user_openai_api_key)

11
backend/repository/chat/create_chat.py
View File

@ -1,8 +1,9 @@
from logger import get_logger
from models.settings import common_dependencies
from dataclasses import dataclass
from models.chat import Chat
from uuid import UUID
from logger import get_logger
from models.chat import Chat
from models.settings import common_dependencies
logger = get_logger(__name__)
@ -15,7 +16,7 @@ class CreateChatProperties:
self.name = name
def create_chat(user_id: str, chat_data: CreateChatProperties) -> Chat:
def create_chat(user_id: UUID, chat_data: CreateChatProperties) -> Chat:
commons = common_dependencies()
# Chat is created upon the user's first question asked
@ -23,7 +24,7 @@ def create_chat(user_id: str, chat_data: CreateChatProperties) -> Chat:
# Insert a new row into the chats table
new_chat = {
"user_id": user_id,
"user_id": str(user_id),
"chat_name": chat_data.name,
}
insert_response = commons["supabase"].table("chats").insert(new_chat).execute()

6
backend/repository/chat/get_user_chats.py
View File

@ -1,8 +1,10 @@
from models.settings import common_dependencies
from typing import List
from models.chat import Chat
from models.settings import common_dependencies
def get_user_chats(user_id: str) -> list[Chat]:
def get_user_chats(user_id: str) -> List[Chat]:
commons = common_dependencies()
response = (
commons["supabase"]

44
backend/routes/api_key_routes.py
View File

@ -10,7 +10,6 @@ from logger import get_logger
from models.settings import CommonsDep
from models.users import User
from pydantic import BaseModel
from utils.users import fetch_user_id_from_credentials
logger = get_logger(__name__)
@ -46,8 +45,6 @@ async def create_api_key(
the user. It returns the newly created API key.
"""
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
new_key_id = str(uuid4())
new_api_key = token_hex(16)
api_key_inserted = False
@ -55,19 +52,13 @@ async def create_api_key(
while not api_key_inserted:
try:
# Attempt to insert new API key into database
commons["supabase"].table("api_keys").insert(
[
{
"key_id": new_key_id,
"user_id": user_id,
"api_key": new_api_key,
"creation_time": datetime.utcnow().strftime(
"%Y-%m-%d %H:%M:%S"
),
"is_active": True,
}
]
).execute()
commons['supabase'].table('api_keys').insert([{
"key_id": new_key_id,
"user_id": current_user.id,
"api_key": new_api_key,
"creation_time": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
"is_active": True
}]).execute()
api_key_inserted = True
@ -96,12 +87,10 @@ async def delete_api_key(
"""
commons["supabase"].table("api_keys").update(
{
"is_active": False,
"deleted_time": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
}
).match({"key_id": key_id, "user_id": current_user.user_id}).execute()
commons['supabase'].table('api_keys').update({
"is_active": False,
"deleted_time": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")
}).match({"key_id": key_id, "user_id": current_user.id}).execute()
return {"message": "API key deleted."}
@ -125,14 +114,5 @@ async def get_api_keys(
containing the key ID and creation time for each API key.
"""
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
response = (
commons["supabase"]
.table("api_keys")
.select("key_id, creation_time")
.filter("user_id", "eq", user_id)
.filter("is_active", "eq", True)
.execute()
)
response = commons['supabase'].table('api_keys').select("key_id, creation_time").filter('user_id', 'eq', current_user.id).filter('is_active', 'eq', True).execute()
return response.data

56
backend/routes/brain_routes.py
View File

@ -4,11 +4,10 @@ from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Request
from logger import get_logger
from models.brains import Brain
from models.brains import Brain, get_default_user_brain
from models.settings import common_dependencies
from models.users import User
from pydantic import BaseModel
from utils.users import fetch_user_id_from_credentials
logger = get_logger(__name__)
@ -37,13 +36,26 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)):
This endpoint retrieves all the brains associated with the current authenticated user. It returns a list of brains objects
containing the brain ID and brain name for each brain.
"""
commons = common_dependencies()
brain = Brain()
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
brains = brain.get_user_brains(user_id)
brains = brain.get_user_brains(current_user.id)
return {"brains": brains}
@brain_router.get("/brains/default", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def get_default_brain_endpoint(current_user: User = Depends(get_current_user)):
"""
Retrieve the default brain for the current user.
- `current_user`: The current authenticated user.
- Returns the default brain for the user.
This endpoint retrieves the default brain associated with the current authenticated user.
The default brain is defined as the brain marked as default in the brains_users table.
"""
default_brain = get_default_user_brain(current_user)
return default_brain
# get one brain
@brain_router.get(
"/brains/{brain_id}", dependencies=[Depends(AuthBearer())], tags=["Brain"]
@ -58,7 +70,7 @@ async def get_brain_endpoint(brain_id: UUID):
This endpoint retrieves the details of a specific brain identified by the provided brain ID. It returns the brain ID and its
history, which includes the brain messages exchanged in the brain.
"""
brain = Brain(brain_id=brain_id)
brain = Brain(id=brain_id)
brains = brain.get_brain_details()
if len(brains) > 0:
return {
@ -74,12 +86,17 @@ async def get_brain_endpoint(brain_id: UUID):
@brain_router.delete(
"/brains/{brain_id}", dependencies=[Depends(AuthBearer())], tags=["Brain"]
)
async def delete_brain_endpoint(brain_id: UUID):
async def delete_brain_endpoint(brain_id: UUID, current_user: User = Depends(get_current_user),):
"""
Delete a specific brain by brain ID.
"""
brain = Brain(brain_id=brain_id)
# [TODO] check if the user is the owner of the brain
current_user.id,
brain = Brain(id=brain_id)
brain.delete_brain()
return {"message": f"{brain_id} has been deleted."}
@ -95,8 +112,7 @@ class BrainObject(BaseModel):
# create new brain
@brain_router.post("/brains", dependencies=[Depends(AuthBearer())], tags=["Brain"])
async def create_brain_endpoint(
request: Request,
async def brain_endpoint(
brain: BrainObject,
current_user: User = Depends(get_current_user),
):
@ -109,15 +125,17 @@ async def create_brain_endpoint(
temperature
In the brains table & in the brains_users table and put the creator user as 'Owner'
"""
commons = common_dependencies()
brain = Brain(name=brain.name)
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
created_brain = brain.create_brain(brain.name)[0]
# create a brain X user entry
brain.create_brain_user(created_brain["brain_id"], user_id, rights="Owner")
return {"id": created_brain["brain_id"], "name": created_brain["name"]}
brain.create_brain()
default_brain = get_default_user_brain(current_user)
if default_brain:
# create a brain X user entry
brain.create_brain_user(user_id = current_user.id, rights="Owner", default_brain=False)
else:
brain.create_brain_user(user_id = current_user.id, rights="Owner", default_brain=True)
return {"id": brain.id, "name": brain.name}
# update existing brain
@brain_router.put(
@ -139,7 +157,7 @@ async def update_brain_endpoint(
Return modified brain ? No need -> do an optimistic update
"""
commons = common_dependencies()
brain = Brain(brain_id=brain_id)
brain = Brain(id=brain_id)
# Add new file to brain , il file_sha1 already exists in brains_vectors -> out (not now)
if brain.file_sha1:

70
backend/routes/chat_routes.py
View File

@ -1,13 +1,14 @@
import os
import time
from http.client import HTTPException
from typing import List
from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, HTTPException, Request
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from llm.brainpicking import BrainPicking
from llm.BrainPickingOpenAIFunctions.BrainPickingOpenAIFunctions import (
BrainPickingOpenAIFunctions,
)
from llm.BrainPickingOpenAIFunctions.BrainPickingOpenAIFunctions import \
BrainPickingOpenAIFunctions
from llm.PrivateBrainPicking import PrivateBrainPicking
from models.chat import Chat, ChatHistory
from models.chats import ChatQuestion
@ -19,7 +20,6 @@ from repository.chat.get_chat_history import get_chat_history
from repository.chat.get_user_chats import get_user_chats
from repository.chat.update_chat import ChatUpdatableProperties, update_chat
from repository.chat.update_chat_history import update_chat_history
from utils.users import fetch_user_id_from_credentials, update_user_request_count
chat_router = APIRouter()
@ -39,19 +39,6 @@ def delete_chat_from_db(commons, chat_id):
commons["supabase"].table("chats").delete().match({"chat_id": chat_id}).execute()
def fetch_user_stats(commons, user, date):
response = (
commons["supabase"]
.from_("users")
.select("*")
.filter("email", "eq", user.email)
.filter("date", "eq", date)
.execute()
)
userItem = next(iter(response.data or []), {"requests_count": 0})
return userItem
# get all chats
@chat_router.get("/chat", dependencies=[Depends(AuthBearer())], tags=["Chat"])
async def get_chats(current_user: User = Depends(get_current_user)):
@ -65,8 +52,7 @@ async def get_chats(current_user: User = Depends(get_current_user)):
containing the chat ID and chat name for each chat.
"""
commons = common_dependencies()
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
chats = get_user_chats(user_id)
chats = get_user_chats( current_user.id)
return {"chats": chats}
@ -97,9 +83,9 @@ async def update_chat_metadata_handler(
"""
commons = common_dependencies()
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
chat = get_chat_by_id(chat_id)
if user_id != chat.user_id:
if current_user.id != chat.user_id:
raise HTTPException(
status_code=403, detail="You should be the owner of the chat to update it."
)
@ -108,20 +94,15 @@ async def update_chat_metadata_handler(
# helper method for update and create chat
def check_user_limit(
email,
user_openai_api_key: str = None,
user : User,
):
if user_openai_api_key is None:
if user.user_openai_api_key is None:
date = time.strftime("%Y%m%d")
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
commons = common_dependencies()
userItem = fetch_user_stats(commons, User(email=email), date)
old_request_count = userItem["requests_count"]
update_user_request_count(
commons, email, date, requests_count=old_request_count + 1
)
if old_request_count >= float(max_requests_number):
user.increment_user_request_count( date )
if user.requests_count >= float(max_requests_number):
raise HTTPException(
status_code=429,
detail="You have reached the maximum number of requests for today.",
@ -140,9 +121,7 @@ async def create_chat_handler(
Create a new chat with initial chat messages.
"""
commons = common_dependencies()
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
return create_chat(user_id=user_id, chat_data=chat_data)
return create_chat(user_id=current_user.id,chat_data=chat_data)
# add new question to chat
@ -153,11 +132,13 @@ async def create_question_handler(
request: Request,
chat_question: ChatQuestion,
chat_id: UUID,
brain_id: UUID = Query(..., description="The ID of the brain"),
current_user: User = Depends(get_current_user),
) -> ChatHistory:
current_user.user_openai_api_key = request.headers.get("Openai-Api-Key")
print("current_user", current_user)
try:
user_openai_api_key = request.headers.get("Openai-Api-Key")
check_user_limit(current_user.email, user_openai_api_key)
check_user_limit(current_user)
llm_settings = LLMSettings()
openai_function_compatible_models = [
"gpt-3.5-turbo-0613",
@ -169,8 +150,8 @@ async def create_question_handler(
chat_id=str(chat_id),
temperature=chat_question.temperature,
max_tokens=chat_question.max_tokens,
user_id=current_user.email,
user_openai_api_key=user_openai_api_key,
brain_id = brain_id,
user_openai_api_key=current_user.user_openai_api_key,
)
answer = gpt_answer_generator.generate_answer(chat_question.question)
elif chat_question.model in openai_function_compatible_models:
@ -181,8 +162,8 @@ async def create_question_handler(
temperature=chat_question.temperature,
max_tokens=chat_question.max_tokens,
# TODO: use user_id in vectors table instead of email
user_email=current_user.email,
user_openai_api_key=user_openai_api_key,
brain_id = brain_id,
user_openai_api_key=current_user.user_openai_api_key,
)
answer = gpt_answer_generator.generate_answer(chat_question.question)
else:
@ -191,8 +172,8 @@ async def create_question_handler(
model=chat_question.model,
max_tokens=chat_question.max_tokens,
temperature=chat_question.temperature,
user_id=current_user.email,
user_openai_api_key=user_openai_api_key,
brain_id = brain_id,
user_openai_api_key=current_user.user_openai_api_key,
)
answer = brainPicking.generate_answer(chat_question.question)
@ -212,7 +193,6 @@ async def create_question_handler(
)
async def get_chat_history_handler(
chat_id: UUID,
current_user: User = Depends(get_current_user),
) -> list[ChatHistory]:
) -> List[ChatHistory]:
# TODO: RBAC with current_user
return get_chat_history(chat_id)

46
backend/routes/crawl_routes.py
View File

@ -1,11 +1,14 @@
import os
import shutil
from tempfile import SpooledTemporaryFile
from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from crawl.crawler import CrawlWebsite
from fastapi import APIRouter, Depends, Request, UploadFile
from models.settings import CommonsDep, common_dependencies
from fastapi import APIRouter, Depends, Query, Request, UploadFile
from models.brains import Brain
from models.files import File
from models.settings import common_dependencies
from models.users import User
from parsers.github import process_github
from utils.file import convert_bytes
@ -13,35 +16,22 @@ from utils.processors import filter_file
crawl_router = APIRouter()
def get_unique_user_data(commons, user):
"""
Retrieve unique user data vectors.
"""
user_vectors_response = commons['supabase'].table("vectors").select(
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
.filter("user_id", "eq", user.email)\
.execute()
documents = user_vectors_response.data # Access the data from the response
# Convert each dictionary to a tuple of items, then to a set to remove duplicates, and then back to a dictionary
user_unique_vectors = [dict(t) for t in set(tuple(d.items()) for d in documents)]
return user_unique_vectors
@crawl_router.post("/crawl/", dependencies=[Depends(AuthBearer())], tags=["Crawl"])
async def crawl_endpoint(request: Request, crawl_website: CrawlWebsite, enable_summarization: bool = False, current_user: User = Depends(get_current_user)):
async def crawl_endpoint(request: Request, crawl_website: CrawlWebsite, brain_id: UUID = Query(..., description="The ID of the brain"),enable_summarization: bool = False, current_user: User = Depends(get_current_user)):
"""
Crawl a website and process the crawled data.
"""
# [TODO] check if the user is the owner/editor of the brain
brain = Brain(id= brain_id)
commons = common_dependencies()
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
if request.headers.get('Openai-Api-Key'):
max_brain_size = os.getenv("MAX_BRAIN_SIZE_WITH_KEY",209715200)
user_unique_vectors = get_unique_user_data(commons, current_user)
current_brain_size = sum(float(doc['size']) for doc in user_unique_vectors)
brain.max_brain_size = os.getenv("MAX_BRAIN_SIZE_WITH_KEY",209715200)
file_size = 1000000
remaining_free_space = float(max_brain_size) - (current_brain_size)
remaining_free_space = brain.remaining_brain_size
if remaining_free_space - file_size < 0:
message = {"message": f"❌ User's brain will exceed maximum capacity with this upload. Maximum file allowed is : {convert_bytes(remaining_free_space)}", "type": "error"}
@ -54,8 +44,12 @@ async def crawl_endpoint(request: Request, crawl_website: CrawlWebsite, enable_s
shutil.copyfileobj(f, spooled_file)
# Pass the SpooledTemporaryFile to UploadFile
file = UploadFile(file=spooled_file, filename=file_name)
message = await filter_file(commons, file, enable_summarization, user=current_user, openai_api_key=request.headers.get('Openai-Api-Key', None))
uploadFile = UploadFile(file=spooled_file, filename=file_name)
file = File(file = uploadFile)
# check remaining free space here !!
message = await filter_file(commons, file, enable_summarization, brain.id, openai_api_key=request.headers.get('Openai-Api-Key', None))
return message
else:
message = await process_github(commons,crawl_website.url, "false", user=current_user, supabase=commons['supabase'], user_openai_api_key=request.headers.get('Openai-Api-Key', None))
# check remaining free space here !!
message = await process_github(commons,crawl_website.url, "false", brain_id, user_openai_api_key=request.headers.get('Openai-Api-Key', None))

51
backend/routes/explore_routes.py
View File

@ -1,35 +1,22 @@
from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends
from models.settings import CommonsDep, common_dependencies
from fastapi import APIRouter, Depends, Query
from models.brains import Brain
from models.settings import common_dependencies
from models.users import User
explore_router = APIRouter()
def get_unique_user_data(commons, user):
"""
Retrieve unique user data vectors.
"""
response = (
commons["supabase"]
.table("vectors")
.select("name:metadata->>file_name, size:metadata->>file_size", count="exact")
.filter("user_id", "eq", user.email)
.execute()
)
documents = response.data # Access the data from the response
# Convert each dictionary to a tuple of items, then to a set to remove duplicates, and then back to a dictionary
unique_data = [dict(t) for t in set(tuple(d.items()) for d in documents)]
return unique_data
@explore_router.get("/explore", dependencies=[Depends(AuthBearer())], tags=["Explore"])
async def explore_endpoint(current_user: User = Depends(get_current_user)):
async def explore_endpoint(brain_id: UUID = Query(..., description="The ID of the brain"),current_user: User = Depends(get_current_user)):
"""
Retrieve and explore unique user data vectors.
"""
commons = common_dependencies()
unique_data = get_unique_user_data(commons, current_user)
brain = Brain(id=brain_id)
unique_data = brain.get_unique_brain_files()
unique_data.sort(key=lambda x: int(x["size"]), reverse=True)
return {"documents": unique_data}
@ -37,20 +24,14 @@ async def explore_endpoint(current_user: User = Depends(get_current_user)):
@explore_router.delete(
"/explore/{file_name}", dependencies=[Depends(AuthBearer())], tags=["Explore"]
)
async def delete_endpoint(file_name: str, credentials: dict = Depends(AuthBearer())):
async def delete_endpoint(file_name: str, current_user: User = Depends(get_current_user), brain_id: UUID = Query(..., description="The ID of the brain")):
"""
Delete a specific user file by file name.
"""
commons = common_dependencies()
user = User(email=credentials.get("email", "none"))
# Cascade delete the summary from the database first, because it has a foreign key constraint
commons["supabase"].table("summaries").delete().match(
{"metadata->>file_name": file_name}
).execute()
commons["supabase"].table("vectors").delete().match(
{"metadata->>file_name": file_name, "user_id": user.email}
).execute()
return {"message": f"{file_name} of user {user.email} has been deleted."}
brain = Brain(id=brain_id)
brain.delete_file_from_brain(file_name)
return {"message": f"{file_name} of brain {brain_id} has been deleted by user {current_user.email}."}
@explore_router.get(
@ -62,6 +43,8 @@ async def download_endpoint(
"""
Download a specific user file by file name.
"""
# check if user has the right to get the file: add brain_id to the query
commons = common_dependencies()
response = (
commons["supabase"]
@ -70,7 +53,7 @@ async def download_endpoint(
"metadata->>file_name, metadata->>file_size, metadata->>file_extension, metadata->>file_url",
"content",
)
.match({"metadata->>file_name": file_name, "user_id": current_user.email})
.match({"metadata->>file_name": file_name})
.execute()
)
documents = response.data

20
backend/routes/stream_routes.py
View File

@ -1,9 +1,10 @@
import asyncio
import os
from typing import AsyncIterable, Awaitable
from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends
from auth.auth_bearer import AuthBearer
from fastapi import APIRouter, Depends, Query
from fastapi.responses import StreamingResponse
from langchain.callbacks import AsyncIteratorCallbackHandler
from langchain.chains import ConversationalRetrievalChain
@ -15,11 +16,10 @@ from llm.prompt.CONDENSE_PROMPT import CONDENSE_QUESTION_PROMPT
from logger import get_logger
from models.chats import ChatMessage
from models.settings import CommonsDep, common_dependencies
from models.users import User
from supabase import create_client
from utils.users import fetch_user_id_from_credentials
from vectorstore.supabase import CustomSupabaseVectorStore
from supabase import create_client
logger = get_logger(__name__)
stream_router = APIRouter()
@ -65,15 +65,13 @@ async def send_message(
await task
def create_chain(commons: CommonsDep, current_user: User):
user_id = fetch_user_id_from_credentials(commons, {"email": current_user.email})
def create_chain(commons: CommonsDep, brain_id: UUID):
embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key)
supabase_client = create_client(supabase_url, supabase_service_key)
vector_store = CustomSupabaseVectorStore(
supabase_client, embeddings, table_name="vectors", user_id=user_id
supabase_client, embeddings, table_name="vectors", brain_id=brain_id
)
generator_llm = ChatOpenAI(
@ -111,11 +109,11 @@ def create_chain(commons: CommonsDep, current_user: User):
@stream_router.post("/stream", dependencies=[Depends(AuthBearer())], tags=["Stream"])
async def stream(
chat_message: ChatMessage,
current_user: User = Depends(get_current_user),
brain_id: UUID = Query(..., description="The ID of the brain"),
) -> StreamingResponse:
commons = common_dependencies()
qa_chain, callback = create_chain(commons, current_user)
qa_chain, callback = create_chain(commons, brain_id)
return StreamingResponse(
send_message(chat_message, qa_chain, callback),

37
backend/routes/upload_routes.py
View File

@ -1,8 +1,11 @@
import os
from uuid import UUID
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Request, UploadFile
from models.settings import CommonsDep, common_dependencies
from fastapi import APIRouter, Depends, Query, Request, UploadFile
from models.brains import Brain
from models.files import File
from models.settings import common_dependencies
from models.users import User
from utils.file import convert_bytes, get_file_size
from utils.processors import filter_file
@ -19,12 +22,9 @@ def get_user_vectors(commons, user):
user_unique_vectors = [dict(t) for t in set(tuple(d.items()) for d in documents)]
return user_unique_vectors
def calculate_remaining_space(request, max_brain_size, max_brain_size_with_own_key, current_brain_size):
remaining_free_space = float(max_brain_size_with_own_key) - current_brain_size if request.headers.get('Openai-Api-Key') else float(max_brain_size) - current_brain_size
return remaining_free_space
@upload_router.post("/upload", dependencies=[Depends(AuthBearer())], tags=["Upload"])
async def upload_file(request: Request, file: UploadFile, enable_summarization: bool = False, current_user: User = Depends(get_current_user)):
async def upload_file(request: Request, uploadFile: UploadFile, brain_id: UUID = Query(..., description="The ID of the brain"), enable_summarization: bool = False, current_user: User = Depends(get_current_user)):
"""
Upload a file to the user's storage.
@ -37,20 +37,27 @@ async def upload_file(request: Request, file: UploadFile, enable_summarization:
and ensures that the file size does not exceed the maximum capacity. If the file is within the allowed size limit,
it can optionally apply summarization to the file's content. The response message will indicate the status of the upload.
"""
print(brain_id,"brain_id")
# [TODO] check if the user is the owner/editor of the brain
brain = Brain(id = brain_id)
print("brain", brain)
commons = common_dependencies()
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200)
if request.headers.get('Openai-Api-Key'):
brain.max_brain_size = os.getenv("MAX_BRAIN_SIZE_WITH_KEY",209715200)
remaining_free_space = brain.remaining_brain_size
user_unique_vectors = get_user_vectors(commons, current_user)
current_brain_size = sum(float(doc['size']) for doc in user_unique_vectors)
remaining_free_space = calculate_remaining_space(request, max_brain_size, max_brain_size_with_own_key, current_brain_size)
file_size = get_file_size(file)
file_size = get_file_size(uploadFile)
file = File(file=uploadFile)
if remaining_free_space - file_size < 0:
message = {"message": f"❌ User's brain will exceed maximum capacity with this upload. Maximum file allowed is : {convert_bytes(remaining_free_space)}", "type": "error"}
else:
message = await filter_file(commons, file, enable_summarization, current_user, openai_api_key=request.headers.get('Openai-Api-Key', None))
message = await filter_file(commons, file, enable_summarization, brain_id=brain_id ,openai_api_key=request.headers.get('Openai-Api-Key', None))
return message

21
backend/routes/user_routes.py
View File

@ -3,7 +3,7 @@ import time
from auth.auth_bearer import AuthBearer, get_current_user
from fastapi import APIRouter, Depends, Request
from models.settings import CommonsDep, common_dependencies
from models.brains import Brain, get_default_user_brain
from models.users import User
user_router = APIRouter()
@ -25,11 +25,6 @@ def get_user_vectors(commons, email):
return user_vectors_response.data
def get_user_request_stats(commons, email):
requests_stats = commons['supabase'].from_('users').select(
'*').filter("email", "eq", email).execute()
return requests_stats.data
@user_router.get("/user", dependencies=[Depends(AuthBearer())], tags=["User"])
async def get_user_endpoint(request: Request, current_user: User = Depends(get_current_user)):
@ -44,11 +39,10 @@ async def get_user_endpoint(request: Request, current_user: User = Depends(get_c
user's uploaded vectors, and the maximum brain size is obtained from the environment variables. The requests statistics provide
information about the user's API usage.
"""
commons = common_dependencies()
user_vectors = get_user_vectors(commons, current_user.email)
user_unique_vectors = get_unique_documents(user_vectors)
# user_vectors = get_user_vectors(commons, current_user.email)
# user_unique_vectors = get_unique_documents(user_vectors)
current_brain_size = sum(float(doc.get('size', 0)) for doc in user_unique_vectors)
# current_brain_size = sum(float(doc.get('size', 0)) for doc in user_unique_vectors)
max_brain_size = int(os.getenv("MAX_BRAIN_SIZE", 0))
if request.headers.get('Openai-Api-Key'):
@ -56,12 +50,13 @@ async def get_user_endpoint(request: Request, current_user: User = Depends(get_c
date = time.strftime("%Y%m%d")
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
requests_stats = get_user_request_stats(commons, current_user.email)
requests_stats = current_user.get_user_request_stats()
defaultBrain = Brain(id=get_default_user_brain(current_user)['id'])
return {"email": current_user.email,
"max_brain_size": max_brain_size,
"current_brain_size": current_brain_size,
"current_brain_size": defaultBrain.brain_size,
"max_requests_number": max_requests_number,
"requests_stats" : requests_stats,
"date": date,

28
backend/utils/processors.py
View File

@ -1,10 +1,7 @@
import os
from fastapi import Depends, FastAPI, UploadFile
from models.files import File
from models.settings import CommonsDep
from models.users import User
from parsers.audio import process_audio
from parsers.common import file_already_exists
from parsers.csv import process_csv
from parsers.docx import process_docx
from parsers.epub import process_epub
@ -15,7 +12,6 @@ from parsers.odt import process_odt
from parsers.pdf import process_pdf
from parsers.powerpoint import process_powerpoint
from parsers.txt import process_txt
from supabase import Client
file_processors = {
".txt": process_txt,
@ -41,16 +37,18 @@ file_processors = {
async def filter_file(commons: CommonsDep, file: UploadFile, enable_summarization: bool, user: User, openai_api_key):
if await file_already_exists(commons['supabase'], file, user):
return {"message": f"🤔 {file.filename} already exists.", "type": "warning"}
elif file.file._file.tell() < 1:
return {"message": f"❌ {file.filename} is empty.", "type": "error"}
async def filter_file(commons: CommonsDep, file: File, enable_summarization: bool, brain_id, openai_api_key):
await file.compute_file_sha1()
print("file sha1", file.file_sha1)
if file.file_already_exists( brain_id):
return {"message": f"🤔 {file.file.filename} already exists in brain {brain_id}.", "type": "warning"}
elif file.file_is_empty():
return {"message": f"❌ {file.file.filename} is empty.", "type": "error"}
else:
file_extension = os.path.splitext(file.filename)[-1].lower() # Convert file extension to lowercase
if file_extension in file_processors:
await file_processors[file_extension](commons,file, enable_summarization, user ,openai_api_key )
return {"message": f"✅ {file.filename} has been uploaded.", "type": "success"}
if file.file_extension in file_processors:
await file_processors[file.file_extension](commons,file, enable_summarization, brain_id ,openai_api_key )
return {"message": f"✅ {file.file.filename} has been uploaded to brain {brain_id}.", "type": "success"}
else:
return {"message": f"❌ {file.filename} is not supported.", "type": "error"}
return {"message": f"❌ {file.file.filename} is not supported.", "type": "error"}

44
backend/utils/users.py
View File

@ -1,4 +1,3 @@
import time
from logger import get_logger
from models.settings import CommonsDep
@ -6,45 +5,10 @@ from models.users import User
logger = get_logger(__name__)
def create_user(commons: CommonsDep, user:User, date):
logger.info(f"New user entry in db document for user {user.email}")
def create_user(commons: CommonsDep, email, date):
logger.info(f"New user entry in db document for user {email}")
return (
commons["supabase"]
.table("users")
.insert({"email": email, "date": date, "requests_count": 1})
.execute()
)
return(commons['supabase'].table("users").insert(
{"user_id": user.id, "email": user.email, "date": date, "requests_count": 1}).execute())
def update_user_request_count(commons: CommonsDep, email, date, requests_count):
logger.info(f"User {email} request count updated to {requests_count}")
commons["supabase"].table("users").update({"requests_count": requests_count}).match(
{"email": email, "date": date}
).execute()
def fetch_user_id_from_credentials(commons: CommonsDep, credentials):
user = User(email=credentials.get("email", "none"))
# Fetch the user's UUID based on their email
response = (
commons["supabase"]
.from_("users")
.select("user_id")
.filter("email", "eq", user.email)
.execute()
)
userItem = next(iter(response.data or []), {})
if userItem == {}:
date = time.strftime("%Y%m%d")
create_user_response = create_user(commons, email=user.email, date=date)
user_id = create_user_response.data[0]["user_id"]
else:
user_id = userItem["user_id"]
return user_id

19
backend/utils/vectors.py
View File

@ -1,9 +1,8 @@
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.schema import Document
from llm.brainpicking import BrainPicking, BrainSettings
from llm.summarization import llm_evaluate_summaries, llm_summerize
from llm.brainpicking import BrainSettings
from llm.summarization import llm_summerize
from logger import get_logger
from models.chats import ChatMessage
from models.settings import BrainSettings, CommonsDep
from pydantic import BaseModel
@ -13,8 +12,8 @@ logger = get_logger(__name__)
class Neurons(BaseModel):
commons: CommonsDep
settings = BrainSettings()
def create_vector(self, user_id, doc, user_openai_api_key=None):
def create_vector(self, doc, user_openai_api_key=None):
logger.info(f"Creating vector for document")
logger.info(f"Document: {doc}")
if user_openai_api_key:
@ -24,9 +23,8 @@ class Neurons(BaseModel):
try:
sids = self.commons["documents_vector_store"].add_documents([doc])
if sids and len(sids) > 0:
self.commons["supabase"].table("vectors").update(
{"user_id": user_id}
).match({"id": sids[0]}).execute()
return sids
except Exception as e:
logger.error(f"Error creating vector for document {e}")
@ -58,6 +56,5 @@ def create_summary(commons: CommonsDep, document_id, content, metadata):
summary_doc_with_metadata = Document(page_content=summary, metadata=metadata)
sids = commons["summaries_vector_store"].add_documents([summary_doc_with_metadata])
if sids and len(sids) > 0:
commons["supabase"].table("summaries").update(
{"document_id": document_id}
).match({"id": sids[0]}).execute()
commons['supabase'].table("summaries").update(
{"document_id": document_id}).match({"id": sids[0]}).execute()

9
backend/vectorstore/supabase.py
View File

@ -3,23 +3,24 @@ from typing import Any, List
from langchain.docstore.document import Document
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import SupabaseVectorStore
from supabase import Client
class CustomSupabaseVectorStore(SupabaseVectorStore):
"""A custom vector store that uses the match_vectors table instead of the vectors table."""
user_id: str
brain_id: str = "none"
def __init__(
self,
client: Client,
embedding: OpenAIEmbeddings,
table_name: str,
user_id: str,
brain_id: str = "none",
):
super().__init__(client, embedding, table_name)
self.user_id = user_id
self.brain_id = brain_id
def similarity_search(
self,
@ -36,7 +37,7 @@ class CustomSupabaseVectorStore(SupabaseVectorStore):
{
"query_embedding": query_embedding,
"match_count": k,
"p_user_id": self.user_id,
"p_brain_id": str(self.brain_id),
},
).execute()

9
frontend/app/chat/[chatId]/hooks/useChatService.ts
View File

@ -1,3 +1,4 @@
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { useAxios } from "@/lib/hooks";
import { ChatEntity, ChatHistory, ChatQuestion } from "../types";
@ -5,7 +6,7 @@ import { ChatEntity, ChatHistory, ChatQuestion } from "../types";
// eslint-disable-next-line @typescript-eslint/explicit-module-boundary-types
export const useChatService = () => {
const { axiosInstance } = useAxios();
const { currentBrain } = useBrainContext();
const createChat = async ({ name }: { name: string }) => {
return axiosInstance.post<ChatEntity>(`/chat`, { name });
};
@ -24,9 +25,13 @@ export const useChatService = () => {
chatId: string,
chatQuestion: ChatQuestion
): Promise<ChatHistory> => {
if (currentBrain?.id === undefined) {
throw new Error("No current brain");
}
return (
await axiosInstance.post<ChatHistory>(
`/chat/${chatId}/question`,
`/chat/${chatId}/question/?brain_id=${currentBrain.id}`,
chatQuestion
)
).data;

16
frontend/app/explore/DocumentItem/index.tsx
View File

@ -13,11 +13,11 @@ import { AnimatedCard } from "@/lib/components/ui/Card";
import Ellipsis from "@/lib/components/ui/Ellipsis";
import Modal from "@/lib/components/ui/Modal";
import { useSupabase } from "@/lib/context/SupabaseProvider";
import { useAxios } from "@/lib/hooks";
import { useToast } from "@/lib/hooks/useToast";
import { useAxios, useToast } from "@/lib/hooks";
import { Document } from "@/lib/types/Document";
import { useEventTracking } from "@/services/analytics/useEventTracking";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import DocumentData from "./DocumentData";
interface DocumentProps {
@ -32,6 +32,7 @@ const DocumentItem = forwardRef(
const { session } = useSupabase();
const { axiosInstance } = useAxios();
const { track } = useEventTracking();
const { currentBrain } = useBrainContext();
if (!session) {
throw new Error("User session not found");
@ -41,9 +42,16 @@ const DocumentItem = forwardRef(
setIsDeleting(true);
void track("DELETE_DOCUMENT");
try {
await axiosInstance.delete(`/explore/${name}`);
if (currentBrain?.id === undefined)
throw new Error("Brain id not found");
await axiosInstance.delete(
`/explore/${name}/?brain_id=${currentBrain.id}`
);
setDocuments((docs) => docs.filter((doc) => doc.name !== name)); // Optimistic update
publish({ variant: "success", text: `${name} deleted.` });
publish({
variant: "success",
text: `${name} deleted from brain ${currentBrain.name}.`,
});
} catch (error) {
console.error(`Error deleting ${name}`, error);
}

32
frontend/app/explore/page.tsx
View File

@ -11,6 +11,9 @@ import { useSupabase } from "@/lib/context/SupabaseProvider";
import { useAxios } from "@/lib/hooks";
import { Document } from "@/lib/types/Document";
import { getBrainFromLocalStorage } from "@/lib/context/BrainProvider/helpers/brainLocalStorage";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { UUID } from "crypto";
import DocumentItem from "./DocumentItem";
const ExplorePage = (): JSX.Element => {
@ -18,20 +21,39 @@ const ExplorePage = (): JSX.Element => {
const [isPending, setIsPending] = useState(true);
const { session } = useSupabase();
const { axiosInstance } = useAxios();
const { setActiveBrain, setDefaultBrain, currentBrain, currentBrainId } =
useBrainContext();
const fetchAndSetActiveBrain = async () => {
const storedBrain = getBrainFromLocalStorage();
if (storedBrain) {
setActiveBrain(storedBrain.id);
return storedBrain;
} else {
const defaultBrain = await setDefaultBrain();
return defaultBrain;
}
};
if (session === null) {
redirect("/login");
}
useEffect(() => {
const fetchDocuments = async () => {
const fetchDocuments = async (brainId: UUID | null) => {
setIsPending(true);
await fetchAndSetActiveBrain();
try {
if (brainId === undefined || brainId === null) {
throw new Error("Brain id not found");
}
console.log(
`Fetching documents from ${process.env.NEXT_PUBLIC_BACKEND_URL}/explore`
`Fetching documents from ${process.env.NEXT_PUBLIC_BACKEND_URL}/explore/?brain_id=${brainId}`
);
const response = await axiosInstance.get<{ documents: Document[] }>(
"/explore"
`/explore/?brain_id=${brainId}`
);
setDocuments(response.data.documents);
} catch (error) {
@ -40,8 +62,8 @@ const ExplorePage = (): JSX.Element => {
}
setIsPending(false);
};
fetchDocuments();
}, [session.access_token, axiosInstance]);
fetchDocuments(currentBrainId);
}, [session.access_token, axiosInstance, currentBrainId]);
return (
<main>

91
frontend/app/upload/components/Crawler/hooks/useCrawler.ts
View File

@ -3,11 +3,10 @@ import { redirect } from "next/navigation";
import { useCallback, useRef, useState } from "react";
import { useSupabase } from "@/lib/context/SupabaseProvider";
import { useAxios } from "@/lib/hooks";
import { useToast } from "@/lib/hooks/useToast";
import { useAxios, useToast } from "@/lib/hooks";
import { useEventTracking } from "@/services/analytics/useEventTracking";
import { UUID } from "crypto";
import { isValidUrl } from "../helpers/isValidUrl";
export const useCrawler = () => {
@ -16,56 +15,64 @@ export const useCrawler = () => {
const { session } = useSupabase();
const { publish } = useToast();
const { axiosInstance } = useAxios();
const { track} = useEventTracking();
const { track } = useEventTracking();
if (session === null) {
redirect("/login");
}
const crawlWebsite = useCallback(async () => {
// Validate URL
const url = urlInputRef.current ? urlInputRef.current.value : null;
const crawlWebsite = useCallback(
async (brainId: UUID | undefined) => {
// Validate URL
const url = urlInputRef.current ? urlInputRef.current.value : null;
if (!url || !isValidUrl(url)) {
// Assuming you have a function to validate URLs
void track("URL_INVALID");
publish({
variant: "danger",
text: "Invalid URL",
});
if (!url || !isValidUrl(url)) {
void track("URL_INVALID");
return;
}
publish({
variant: "danger",
text: "Invalid URL",
});
// Configure parameters
const config = {
url: url,
js: false,
depth: 1,
max_pages: 100,
max_time: 60,
};
return;
}
setCrawling(true);
void track("URL_CRAWLED");
// Configure parameters
const config = {
url: url,
js: false,
depth: 1,
max_pages: 100,
max_time: 60,
};
try {
const response = await axiosInstance.post(`/crawl/`, config);
setCrawling(true);
void track("URL_CRAWLED");
publish({
variant: response.data.type,
text: response.data.message,
});
} catch (error: unknown) {
publish({
variant: "danger",
text: "Failed to crawl website: " + JSON.stringify(error),
});
} finally {
setCrawling(false);
}
}, [session.access_token, publish]);
try {
console.log("Crawling website...", brainId);
if (brainId !== undefined) {
const response = await axiosInstance.post(
`/crawl/?brain_id=${brainId}`,
config
);
publish({
variant: response.data.type,
text: response.data.message,
});
}
} catch (error: unknown) {
publish({
variant: "danger",
text: "Failed to crawl website: " + JSON.stringify(error),
});
} finally {
setCrawling(false);
}
},
[session.access_token, publish]
);
return {
isCrawling,

7
frontend/app/upload/components/Crawler/index.tsx
View File

@ -2,11 +2,13 @@
import Button from "@/lib/components/ui/Button";
import Card from "@/lib/components/ui/Card";
import Field from "@/lib/components/ui/Field";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { useCrawler } from "./hooks/useCrawler";
export const Crawler = (): JSX.Element => {
const { urlInputRef, isCrawling, crawlWebsite } = useCrawler();
const { currentBrain } = useBrainContext();
return (
<div className="w-full">
@ -24,7 +26,10 @@ export const Crawler = (): JSX.Element => {
/>
</div>
<div className="flex flex-col items-center justify-center gap-5">
<Button isLoading={isCrawling} onClick={() => void crawlWebsite()}>
<Button
isLoading={isCrawling}
onClick={() => void crawlWebsite(currentBrain?.id)}
>
Crawl
</Button>
</div>

35
frontend/app/upload/components/FileUploader/hooks/useFileUploader.ts
View File

@ -3,31 +3,40 @@ import { redirect } from "next/navigation";
import { useCallback, useState } from "react";
import { FileRejection, useDropzone } from "react-dropzone";
import { useBrainContext } from "@/lib/context/BrainProvider/hooks/useBrainContext";
import { useSupabase } from "@/lib/context/SupabaseProvider";
import { useAxios } from "@/lib/hooks";
import { useToast } from "@/lib/hooks/useToast";
import { useAxios, useToast } from "@/lib/hooks";
import { useEventTracking } from "@/services/analytics/useEventTracking";
import { useFeature } from "@growthbook/growthbook-react";
import { UUID } from "crypto";
export const useFileUploader = () => {
const { track} = useEventTracking();
const { track } = useEventTracking();
const [isPending, setIsPending] = useState(false);
const { publish } = useToast();
const [files, setFiles] = useState<File[]>([]);
const { session } = useSupabase();
const { currentBrain, createBrain } = useBrainContext();
const { axiosInstance } = useAxios();
const shouldUseMultipleBrains = useFeature("multiple-brains").on;
if (session === null) {
redirect("/login");
}
const upload = useCallback(
async (file: File) => {
async (file: File, brainId: UUID) => {
const formData = new FormData();
formData.append("file", file);
formData.append("uploadFile", file);
try {
const response = await axiosInstance.post(`/upload`, formData);
void track("FILE_UPLOADED");
const response = await axiosInstance.post(
`/upload/?brain_id=${brainId}`,
formData
);
track("FILE_UPLOADED");
publish({
variant: response.data.type,
text:
@ -78,10 +87,20 @@ export const useFileUploader = () => {
return;
}
setIsPending(true);
if (currentBrain?.id !== undefined) {
setFiles([]);
await Promise.all(files.map((file) => upload(file, currentBrain?.id)));
}
console.log("Please select or create a brain to upload a file");
await Promise.all(files.map((file) => upload(file)));
if (currentBrain?.id === undefined && shouldUseMultipleBrains !== true) {
const createdBrainId = await createBrain("Default");
createdBrainId
? await Promise.all(files.map((file) => upload(file, createdBrainId)))
: null;
setFiles([]);
}
setFiles([]);
setIsPending(false);
};

76
frontend/lib/api/brains.ts Normal file
View File

@ -0,0 +1,76 @@
import { AxiosInstance } from "axios";
import { UUID } from "crypto";
import { Brain } from "../context/BrainProvider/types";
export const createBrainFromBackend = async (
axiosInstance: AxiosInstance,
name: string
): Promise<Brain | undefined> => {
try {
const createdBrain = (await axiosInstance.post<Brain>(`/brains`, { name }))
.data;
return createdBrain;
} catch (error) {
console.error(`Error creating brain ${name}`, error);
}
};
export const getUserDefaultBrainFromBackend = async (
axiosInstance: AxiosInstance
): Promise<Brain | undefined> => {
try {
const defaultBrain = (await axiosInstance.get<Brain>(`/brains/default`))
.data;
return { id: defaultBrain.id, name: defaultBrain.name };
} catch (error) {
console.error(`Error getting user's default brain`, error);
}
};
export const getBrainFromBE = async (
axiosInstance: AxiosInstance,
brainId: UUID
): Promise<Brain | undefined> => {
try {
const brain = (await axiosInstance.get<Brain>(`/brains/${brainId}`)).data;
return brain;
} catch (error) {
console.error(`Error getting brain ${brainId}`, error);
throw new Error(`Error getting brain ${brainId}`);
}
};
export const deleteBrainFromBE = async (
axiosInstance: AxiosInstance,
brainId: UUID
): Promise<void> => {
try {
(await axiosInstance.delete(`/brain/${brainId}`)).data;
} catch (error) {
console.error(`Error deleting brain ${brainId}`, error);
throw new Error(`Error deleting brain ${brainId}`);
}
};
export const getAllUserBrainsFromBE = async (
axiosInstance: AxiosInstance
): Promise<Brain[] | undefined> => {
try {
const brains = (await axiosInstance.get<{ brains: Brain[] }>(`/brains`))
.data;
console.log("BRAINS", brains);
return brains.brains;
} catch (error) {
console.error(`Error getting brain for current user}`, error);
throw new Error(`Error getting brain for current user`);
}
};

1
frontend/lib/api/index.ts Normal file
View File

@ -0,0 +1 @@
export * from "./brains";

8
frontend/lib/components/NavBar/components/NavItems/index.tsx
View File

@ -1,5 +1,5 @@
"use client";
import { useFeature } from '@growthbook/growthbook-react';
import { useFeature } from "@growthbook/growthbook-react";
import Link from "next/link";
import { Dispatch, HTMLAttributes, SetStateAction } from "react";
import { MdPerson, MdSettings } from "react-icons/md";
@ -24,8 +24,8 @@ export const NavItems = ({
}: NavItemsProps): JSX.Element => {
const { session } = useSupabase();
const isUserLoggedIn = session?.user !== undefined;
const shouldUseMultipleBrains = useFeature('multiple-brains').on;
const shouldUseMultipleBrains = useFeature("multiple-brains").on;
return (
<ul
className={cn(
@ -59,7 +59,7 @@ export const NavItems = ({
<div className="flex sm:flex-1 sm:justify-end flex-col items-center justify-center sm:flex-row gap-5 sm:gap-2">
{isUserLoggedIn && (
<>
{shouldUseMultipleBrains && (<BrainsDropDown />)}
{shouldUseMultipleBrains && <BrainsDropDown />}
<Link aria-label="account" className="" href={"/user"}>
<MdPerson className="text-2xl" />
</Link>

4
frontend/lib/context/BrainProvider/helpers/brainLocalStorage.ts
View File

@ -2,8 +2,8 @@ import { Brain } from "../types";
const BRAIN_LOCAL_STORAGE_KEY = "userBrains";
export const saveBrainInLocalStorage = (updatedConfig: Brain): void => {
localStorage.setItem(BRAIN_LOCAL_STORAGE_KEY, JSON.stringify(updatedConfig));
export const saveBrainInLocalStorage = (brain: Brain): void => {
localStorage.setItem(BRAIN_LOCAL_STORAGE_KEY, JSON.stringify(brain));
};
export const getBrainFromLocalStorage = (): Brain | undefined => {
const persistedBrain = localStorage.getItem(BRAIN_LOCAL_STORAGE_KEY);

160
frontend/lib/context/BrainProvider/hooks/useBrainState.ts
View File

@ -1,80 +1,32 @@
/* eslint-disable max-lines */
import { AxiosInstance } from "axios";
import { UUID } from "crypto";
import { useCallback, useEffect, useState } from "react";
import {
createBrainFromBackend,
deleteBrainFromBE,
getAllUserBrainsFromBE,
getBrainFromBE,
getUserDefaultBrainFromBackend,
} from "@/lib/api";
import { useAxios, useToast } from "@/lib/hooks";
import {
getBrainFromLocalStorage,
saveBrainInLocalStorage,
} from "../helpers/brainLocalStorage";
import { Brain } from "../types";
const createBrainFromBackend = async (
axiosInstance: AxiosInstance,
name: string
): Promise<Brain | undefined> => {
try {
const createdBrain = (await axiosInstance.post<Brain>(`/brains`, { name }))
.data;
return createdBrain;
} catch (error) {
console.error(`Error creating brain ${name}`, error);
}
};
const getBrainFromBE = async (
axiosInstance: AxiosInstance,
brainId: UUID
): Promise<Brain | undefined> => {
try {
const brain = (await axiosInstance.get<Brain>(`/brains/${brainId}`)).data;
return brain;
} catch (error) {
console.error(`Error getting brain ${brainId}`, error);
throw new Error(`Error getting brain ${brainId}`);
}
};
const deleteBrainFromBE = async (
axiosInstance: AxiosInstance,
brainId: UUID
): Promise<void> => {
try {
(await axiosInstance.delete(`/brain/${brainId}`)).data;
} catch (error) {
console.error(`Error deleting brain ${brainId}`, error);
throw new Error(`Error deleting brain ${brainId}`);
}
};
const getAllUserBrainsFromBE = async (
axiosInstance: AxiosInstance
): Promise<Brain[] | undefined> => {
try {
const brains = (await axiosInstance.get<{ brains: Brain[] }>(`/brains`))
.data;
console.log("BRAINS", brains);
return brains.brains;
} catch (error) {
console.error(`Error getting brain for current user}`, error);
throw new Error(`Error getting brain for current user`);
}
};
export interface BrainStateProps {
currentBrain: Brain | undefined;
currentBrainId: UUID | null;
allBrains: Brain[];
createBrain: (name: string) => Promise<void>;
createBrain: (name: string) => Promise<UUID | undefined>;
deleteBrain: (id: UUID) => Promise<void>;
setActiveBrain: (id: UUID) => void;
getBrainWithId: (brainId: UUID) => Promise<Brain>;
fetchAllBrains: () => Promise<void>;
setDefaultBrain: () => Promise<void>;
}
export const useBrainState = (): BrainStateProps => {
@ -86,15 +38,15 @@ export const useBrainState = (): BrainStateProps => {
const currentBrain = allBrains.find((brain) => brain.id === currentBrainId);
const setActiveBrain = (id: UUID) => {
setCurrentBrainId(id);
};
// options: Record<string, string | unknown>;
const createBrain = async (name: string) => {
const createBrain = async (name: string): Promise<UUID | undefined> => {
const createdBrain = await createBrainFromBackend(axiosInstance, name);
if (createdBrain !== undefined) {
setAllBrains((prevBrains) => [...prevBrains, createdBrain]);
saveBrainInLocalStorage(createdBrain);
return createdBrain.id;
} else {
publish({
variant: "danger",
@ -120,32 +72,6 @@ export const useBrainState = (): BrainStateProps => {
return brain;
};
// const addDocumentToBrain = (brainId: UUID, document: Document) => {
// const brains = [...allBrains];
// brains.forEach((brain) => {
// if (brain.id === brainId) {
// brain.documents?.push(document);
// return; // return as there cannot be more than one brain with that id
// }
// });
// //call update brain with document -> need file sha1
// setAllBrains(brains);
// };
// const removeDocumentFromBrain = (brainId: UUID, sha1: string) => {
// const brains = [...allBrains];
// brains.forEach((brain) => {
// if (brain.id === brainId) {
// brain.documents = brain.documents?.filter((doc) => doc.sha1 !== sha1);
// //remove document endpoint here (use the document hook ?)
// return; // return as there cannot be more than one brain with that id
// }
// });
// setAllBrains(brains);
// };
const fetchAllBrains = useCallback(async () => {
try {
console.log("Fetching all brains for a user");
@ -158,19 +84,63 @@ export const useBrainState = (): BrainStateProps => {
}
}, [axiosInstance]);
const setActiveBrain = useCallback((id: UUID) => {
//get brain with id from BE ?
const newActiveBrain = { id, name: "Default Brain" };
// if (newActiveBrain) {
console.log("newActiveBrain", newActiveBrain);
saveBrainInLocalStorage(newActiveBrain);
setCurrentBrainId(id);
console.log("Setting active brain", newActiveBrain);
// } else {
// console.warn(`No brain found with ID ${id}`);
// }
}, []);
const setDefaultBrain = useCallback(async () => {
console.log("Setting default brain");
const defaultBrain = await getUserDefaultBrainFromBackend(axiosInstance);
console.log("defaultBrain", defaultBrain);
if (defaultBrain) {
saveBrainInLocalStorage(defaultBrain);
setActiveBrain(defaultBrain.id);
} else {
console.warn("No brains found");
}
}, [axiosInstance, setActiveBrain]);
const fetchAndSetActiveBrain = useCallback(async () => {
console.log(
"Fetching and setting active brain use effect in useBrainState"
);
const storedBrain = getBrainFromLocalStorage();
if (storedBrain?.id !== undefined) {
console.log("Setting active brain from local storage");
console.log("storedBrain", storedBrain);
setActiveBrain(storedBrain.id);
} else {
console.log("Setting default brain for first time");
await setDefaultBrain();
}
}, [setDefaultBrain, setActiveBrain]);
useEffect(() => {
void fetchAllBrains();
}, [fetchAllBrains]);
console.log("brainId", currentBrainId);
void fetchAndSetActiveBrain();
}, [fetchAllBrains, fetchAndSetActiveBrain, currentBrainId]);
return {
currentBrain,
currentBrainId,
allBrains,
createBrain,
deleteBrain,
setActiveBrain,
// addDocumentToBrain,
// removeDocumentFromBrain,
getBrainWithId,
fetchAllBrains,
setDefaultBrain,
};
};

3
frontend/lib/context/index.tsx
View File

@ -1,3 +1,2 @@
export * from "./BrainProvider";
export * from './FeatureFlagProvider';
export * from "./FeatureFlagProvider";

91
scripts/20230620170840_add_vectors_brains.sql Normal file
View File

@ -0,0 +1,91 @@
-- Add a 'brain_id' column to 'vectors' table if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'vectors'
AND column_name = 'brain_id'
) THEN
ALTER TABLE vectors ADD COLUMN brain_id UUID;
END IF;
END
$$;
-- Copy corresponding 'user_id' from 'users' table to 'brain_id' in 'vectors' where 'email' matches only if 'brain_id' is NULL or not equal to 'user_id'
UPDATE vectors v
SET brain_id = u.user_id
FROM users u
WHERE v.user_id = u.email AND (v.brain_id IS NULL OR v.brain_id != u.user_id);
-- Delete rows in 'vectors' where 'brain_id' is NULL
DELETE FROM vectors
WHERE brain_id IS NULL;
-- Create a new entry in 'brains' table for each unique 'brain_id' in 'vectors', avoiding duplicates
INSERT INTO brains (brain_id, name, status, model, max_tokens, temperature)
SELECT brain_id, 'Default', 'public', 'gpt-3', '2048', 0.7 FROM vectors
ON CONFLICT (brain_id) DO NOTHING;
-- Create entries in 'brains_vectors' for all entries in 'vectors', avoiding duplicates
INSERT INTO brains_vectors (brain_id, vector_id)
SELECT brain_id, id FROM vectors
ON CONFLICT (brain_id, vector_id) DO NOTHING;
ALTER TABLE brains_users DROP CONSTRAINT brains_users_user_id_fkey;
ALTER TABLE brains_users ALTER COLUMN user_id TYPE TEXT USING user_id::TEXT;
-- Add a 'default_brain' column to 'brains_users' table if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'brains_users'
AND column_name = 'default_brain'
) THEN
ALTER TABLE brains_users ADD COLUMN default_brain BOOLEAN DEFAULT false;
END IF;
END
$$;
INSERT INTO brains_users (brain_id, user_id, default_brain)
SELECT brain_id, user_id, true FROM vectors
ON CONFLICT (brain_id, user_id) DO NOTHING;
-- Update 'default_brain' as 'true' for all current brains if it's NULL
UPDATE brains_users SET default_brain = true WHERE brain_id IN (SELECT brain_id FROM vectors) AND default_brain IS NULL;
-- Remove 'user_id' column if it exists
DO $$
BEGIN
IF EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'vectors'
AND column_name = 'user_id'
) THEN
ALTER TABLE vectors DROP COLUMN user_id;
END IF;
END
$$;
-- Remove 'brain_id' column if it exists
DO $$
BEGIN
IF EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'vectors'
AND column_name = 'brain_id'
) THEN
ALTER TABLE vectors DROP COLUMN brain_id;
END IF;
END
$$;

212
scripts/20230620183620_use_supabase_user_id.sql Normal file
View File

@ -0,0 +1,212 @@
-- Add the 'supabase_id' column to the 'users' table if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'users'
AND column_name = 'supabase_id'
) THEN
ALTER TABLE users ADD COLUMN supabase_id UUID;
END IF;
END
$$;
-- Update the 'supabase_id' column with the corresponding 'id' from 'auth.users'
-- Fails if there's no matching email in auth.users
UPDATE users
SET supabase_id = au.id
FROM auth.users au
WHERE users.email = au.email;
-- Create a copy of old users table for safety
-- Fails if 'users_old' table already exists
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM information_schema.tables
WHERE table_name = 'users_old'
) THEN
CREATE TABLE users_old AS TABLE users;
END IF;
END
$$;
-- Drop the old primary key if it exists
DO $$
BEGIN
IF EXISTS (
SELECT 1
FROM pg_constraint
WHERE conname = 'users_pkey'
) THEN
ALTER TABLE users DROP CONSTRAINT users_pkey CASCADE;
END IF;
END
$$;
-- Rename columns if not already renamed
DO $$
BEGIN
IF EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'users'
AND column_name = 'user_id'
) AND NOT EXISTS(
SELECT 1
FROM information_schema.columns
WHERE table_name = 'users'
AND column_name = 'old_user_id'
) THEN
ALTER TABLE users RENAME COLUMN user_id TO old_user_id;
ALTER TABLE users RENAME COLUMN supabase_id TO user_id;
END IF;
END
$$;
-- Create a new primary key with user_id and date if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM pg_constraint
WHERE conname = 'users_pkey'
) THEN
ALTER TABLE users ADD PRIMARY KEY (user_id, date);
END IF;
END
$$;
-- Update the 'chats' table
-- Drop old foreign key constraint if it exists
DO $$
BEGIN
IF EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'chats_user_id_fkey'
) THEN
ALTER TABLE chats DROP CONSTRAINT chats_user_id_fkey;
END IF;
END
$$;
-- Update user_id in chats
-- Fails if there's no matching old_user_id in users
UPDATE chats
SET user_id = u.user_id::uuid
FROM users u
WHERE chats.user_id::uuid = u.old_user_id::uuid;
-- Add new foreign key constraint if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'chats_user_id_fkey'
) THEN
ALTER TABLE chats ADD CONSTRAINT chats_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id);
END IF;
END
$$;
-- Update the 'brains_users' table
-- Add a new 'new_user_id' column to the 'brains_users' table
ALTER TABLE brains_users ADD COLUMN new_user_id UUID;
-- Update 'new_user_id' in the 'brains_users' table based on the 'email' in the 'users' table
UPDATE brains_users bu
SET new_user_id = u.user_id
FROM users u
WHERE bu.user_id = u.email;
-- Once you are sure that 'new_user_id' has been correctly populated, drop the old 'user_id' column
ALTER TABLE brains_users DROP COLUMN user_id;
-- Rename 'new_user_id' column to 'user_id'
ALTER TABLE brains_users RENAME COLUMN new_user_id TO user_id;
-- Delete users with user_id not in supabase auth
DELETE FROM brains_users
WHERE NOT EXISTS (
SELECT 1
FROM auth.users
WHERE brains_users.user_id = auth.users.id
);
-- Drop old foreign key constraint if it exists
DO $$
BEGIN
IF EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'brains_users_user_id_fkey'
) THEN
ALTER TABLE brains_users DROP CONSTRAINT brains_users_user_id_fkey;
END IF;
END
$$;
-- Update user_id in brains_users
-- Fails if there's no matching old_user_id in users
UPDATE brains_users
SET user_id = u.user_id::uuid
FROM users u
WHERE brains_users.user_id::uuid = u.old_user_id::uuid;
-- Add new foreign key constraints if they don't exist
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'brains_users_user_id_fkey'
) THEN
ALTER TABLE brains_users ADD CONSTRAINT brains_users_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id);
--ALTER TABLE brains_users ADD CONSTRAINT brains_users_brain_id_fkey FOREIGN KEY (brain_id) REFERENCES brains (brain_id);
END IF;
END
$$;
-- Update the 'api_keys' table
-- Drop old foreign key constraint if it exists
DO $$
BEGIN
IF EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'api_keys_user_id_fkey'
) THEN
ALTER TABLE api_keys DROP CONSTRAINT api_keys_user_id_fkey;
END IF;
END
$$;
-- Update user_id in api_keys
-- Fails if there's no matching old_user_id in users
UPDATE api_keys
SET user_id = u.user_id::uuid
FROM users u
WHERE api_keys.user_id::uuid = u.old_user_id::uuid;
-- Add new foreign key constraint if it doesn't exist
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1
FROM information_schema.table_constraints
WHERE constraint_name = 'api_keys_user_id_fkey'
) THEN
ALTER TABLE api_keys ADD CONSTRAINT api_keys_user_id_fkey FOREIGN KEY (user_id) REFERENCES auth.users (id);
END IF;
END
$$;
-- Optionally Drop the 'old_user_id' column from the 'users' table
-- Uncomment if you are sure that it is no longer needed.
--ALTER TABLE users DROP COLUMN old_user_id;

38
scripts/20230627151100_update_match_vectors.sql Normal file
View File

@ -0,0 +1,38 @@
-- Migration script
BEGIN;
-- Drop the old function if exists
DROP FUNCTION IF EXISTS match_vectors(VECTOR(1536), INT, TEXT);
-- Create the new function
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(1536), match_count INT, p_brain_id UUID)
RETURNS TABLE(
id BIGINT,
brain_id UUID,
content TEXT,
metadata JSONB,
embedding VECTOR(1536),
similarity FLOAT
) LANGUAGE plpgsql AS $$
#variable_conflict use_column
BEGIN
RETURN QUERY
SELECT
vectors.id,
brains_vectors.brain_id,
vectors.content,
vectors.metadata,
vectors.embedding,
1 - (vectors.embedding <=> query_embedding) AS similarity
FROM
vectors
INNER JOIN
brains_vectors ON vectors.id = brains_vectors.vector_id
WHERE brains_vectors.brain_id = p_brain_id
ORDER BY
vectors.embedding <=> query_embedding
LIMIT match_count;
END;
$$;
COMMIT;

41
scripts/tables.sql
View File

@ -1,15 +1,16 @@
-- Create users table
CREATE TABLE IF NOT EXISTS users(
user_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
user_id UUID REFERENCES auth.users (id),
email TEXT,
date TEXT,
requests_count INT
requests_count INT,
PRIMARY KEY (user_id, date)
);
-- Create chats table
CREATE TABLE IF NOT EXISTS chats(
chat_id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
user_id UUID REFERENCES users(user_id),
user_id UUID REFERENCES auth.users (id),
creation_time TIMESTAMP DEFAULT current_timestamp,
history JSONB,
chat_name TEXT
@ -31,17 +32,16 @@ CREATE EXTENSION IF NOT EXISTS vector;
-- Create vectors table
CREATE TABLE IF NOT EXISTS vectors (
id BIGSERIAL PRIMARY KEY,
user_id TEXT,
content TEXT,
metadata JSONB,
embedding VECTOR(1536)
);
-- Create function to match vectors
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(1536), match_count INT, p_user_id TEXT)
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(1536), match_count INT, p_brain_id UUID)
RETURNS TABLE(
id BIGINT,
user_id TEXT,
brain_id UUID,
content TEXT,
metadata JSONB,
embedding VECTOR(1536),
@ -51,21 +51,24 @@ RETURNS TABLE(
BEGIN
RETURN QUERY
SELECT
id,
user_id,
content,
metadata,
embedding,
vectors.id,
brains_vectors.brain_id,
vectors.content,
vectors.metadata,
vectors.embedding,
1 - (vectors.embedding <=> query_embedding) AS similarity
FROM
vectors
WHERE vectors.user_id = p_user_id
INNER JOIN
brains_vectors ON vectors.id = brains_vectors.vector_id
WHERE brains_vectors.brain_id = p_brain_id
ORDER BY
vectors.embedding <=> query_embedding
LIMIT match_count;
END;
$$;
-- Create stats table
CREATE TABLE IF NOT EXISTS stats (
time TIMESTAMP,
@ -117,7 +120,7 @@ $$;
-- Create api_keys table
CREATE TABLE IF NOT EXISTS api_keys(
key_id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
user_id UUID REFERENCES users(user_id),
user_id UUID REFERENCES auth.users (id),
api_key TEXT UNIQUE,
creation_time TIMESTAMP DEFAULT current_timestamp,
deleted_time TIMESTAMP,
@ -139,17 +142,17 @@ CREATE TABLE IF NOT EXISTS brains_users (
brain_id UUID,
user_id UUID,
rights VARCHAR(255),
default_brain BOOLEAN DEFAULT false,
PRIMARY KEY (brain_id, user_id),
FOREIGN KEY (user_id) REFERENCES Users(user_id),
FOREIGN KEY (brain_id) REFERENCES Brains(brain_id)
FOREIGN KEY (user_id) REFERENCES auth.users (id),
FOREIGN KEY (brain_id) REFERENCES Brains (brain_id)
);
-- Create brains X vectors table
CREATE TABLE IF NOT EXISTS brains_vectors (
brain_id UUID,
vector_id BIGSERIAL,
rights VARCHAR(255),
vector_id BIGINT,
PRIMARY KEY (brain_id, vector_id),
FOREIGN KEY (vector_id) REFERENCES vectors(id),
FOREIGN KEY (brain_id) REFERENCES brains(brain_id)
FOREIGN KEY (vector_id) REFERENCES vectors (id),
FOREIGN KEY (brain_id) REFERENCES brains (brain_id)
);