mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-09-11 14:36:35 +03:00
Feat/multiple brains files (#361)
This commit is contained in:
parent
e79da8e3cd
commit
ccdc5bb7a7
@ -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"))
|
||||
|
@ -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
|
||||
|
@ -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,
|
||||
|
@ -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(
|
||||
|
@ -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)
|
||||
|
@ -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
109
backend/models/files.py
Normal 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
|
||||
|
@ -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
|
||||
|
@ -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)
|
||||
|
@ -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
|
||||
|
@ -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,
|
||||
)
|
||||
|
@ -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)
|
||||
|
@ -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)
|
||||
|
@ -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"}
|
||||
|
||||
|
@ -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):
|
||||
|
@ -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)
|
||||
|
@ -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)
|
||||
|
@ -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)
|
||||
|
@ -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)
|
||||
|
||||
|
@ -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)
|
||||
|
@ -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)
|
||||
|
@ -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()
|
||||
|
@ -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"]
|
||||
|
@ -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
|
||||
|
@ -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:
|
||||
|
@ -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)
|
||||
|
@ -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))
|
||||
|
@ -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
|
||||
|
@ -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),
|
||||
|
@ -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
|
||||
|
@ -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,
|
||||
|
@ -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"}
|
||||
|
||||
|
@ -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
|
||||
|
@ -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()
|
||||
|
@ -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()
|
||||
|
||||
|
@ -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;
|
||||
|
@ -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);
|
||||
}
|
||||
|
@ -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>
|
||||
|
@ -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,
|
||||
|
@ -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>
|
||||
|
@ -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
76
frontend/lib/api/brains.ts
Normal 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
1
frontend/lib/api/index.ts
Normal file
@ -0,0 +1 @@
|
||||
export * from "./brains";
|
@ -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>
|
||||
|
@ -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);
|
||||
|
@ -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,
|
||||
};
|
||||
};
|
||||
|
@ -1,3 +1,2 @@
|
||||
export * from "./BrainProvider";
|
||||
export * from './FeatureFlagProvider';
|
||||
|
||||
export * from "./FeatureFlagProvider";
|
||||
|
91
scripts/20230620170840_add_vectors_brains.sql
Normal file
91
scripts/20230620170840_add_vectors_brains.sql
Normal 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
212
scripts/20230620183620_use_supabase_user_id.sql
Normal 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
38
scripts/20230627151100_update_match_vectors.sql
Normal 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;
|
@ -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)
|
||||
);
|
||||
|
Loading…
Reference in New Issue
Block a user