mirror of
https://github.com/StanGirard/quivr.git
synced 2024-11-23 21:22:35 +03:00
110 lines
3.3 KiB
Python
110 lines
3.3 KiB
Python
from typing import List, Optional, Tuple
|
|
from uuid import UUID
|
|
|
|
from models.settings import CommonsDep, common_dependencies
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class Brain(BaseModel):
|
|
brain_id: Optional[UUID]
|
|
name: Optional[str] = "New 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] = ""
|
|
_commons: Optional[CommonsDep] = None
|
|
|
|
class Config:
|
|
arbitrary_types_allowed = True
|
|
|
|
@property
|
|
def commons(self) -> CommonsDep:
|
|
if not self._commons:
|
|
self.__class__._commons = common_dependencies()
|
|
return self._commons
|
|
|
|
@classmethod
|
|
def create(cls, *args, **kwargs):
|
|
commons = common_dependencies()
|
|
return cls(commons=commons, *args, **kwargs)
|
|
|
|
def get_user_brains(self, user_id):
|
|
response = (
|
|
self.commons["supabase"]
|
|
.from_("brains_users")
|
|
.select("id:brain_id, brains (id: brain_id, name)")
|
|
.filter("user_id", "eq", user_id)
|
|
.execute()
|
|
)
|
|
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)
|
|
.execute()
|
|
)
|
|
return response.data
|
|
|
|
def delete_brain(self):
|
|
self.commons["supabase"].table("brains").delete().match(
|
|
{"brain_id": self.brain_id}
|
|
).execute()
|
|
|
|
@classmethod
|
|
def create_brain(cls, name):
|
|
commons = common_dependencies()
|
|
response = commons["supabase"].table("brains").insert({"name": name}).execute()
|
|
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()
|
|
)
|
|
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})
|
|
.execute()
|
|
)
|
|
return response.data
|
|
|
|
def get_vector_ids_from_file_sha1(self, file_sha1: str):
|
|
vectorsResponse = (
|
|
self.commons["supabase"]
|
|
.table("vectors")
|
|
.select("id")
|
|
.filter("metadata->>file_sha1", "eq", file_sha1)
|
|
.execute()
|
|
)
|
|
return vectorsResponse.data
|
|
|
|
def update_brain_fields(self):
|
|
self.commons["supabase"].table("brains").update({"name": self.name}).match(
|
|
{"brain_id": self.brain_id}
|
|
).execute()
|
|
|
|
def update_brain_with_file(self, file_sha1: str):
|
|
vector_ids = self.get_vector_ids_from_file_sha1(file_sha1)
|
|
for vector_id in vector_ids:
|
|
self.create_brain_vector(vector_id)
|