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

Refactor integration brain and brain service code

Browse Source
This commit is contained in:
Stan Girard 2024-04-10 11:30:22 +02:00
parent 15e5129f7e
commit 2a3bce5017
3 changed files with 21 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
backend/modules/brain/entity/integration_brain.py
View File

@ -10,6 +10,7 @@ class IntegrationType(str, Enum):
SYNC = "sync"
DOC = "doc"
class IntegrationBrainTag(str, Enum):
NEW = "new"
RECOMMENDED = "recommended"

39
backend/modules/brain/service/brain_service.py
View File

@ -59,9 +59,7 @@ class BrainService:
return self.brain_repository.get_brain_by_id(brain_id)
def get_integration_brain(self, brain_id) -> IntegrationEntity | None:
return self.integration_brains_repository.get_integration_brain(
brain_id
)
return self.integration_brains_repository.get_integration_brain(brain_id)
def find_brain_from_question(
self,
@ -105,22 +103,23 @@ class BrainService:
if brain_id_to_use and not brain_to_use:
brain_to_use = self.get_brain_by_id(brain_id_to_use)
# Calculate the closest brains to the question
list_brains = vector_store.find_brain_closest_query(user.id, question)
else:
# Calculate the closest brains to the question
list_brains = vector_store.find_brain_closest_query(user.id, question)
unique_list_brains = []
seen_brain_ids = set()
unique_list_brains = []
seen_brain_ids = set()
for brain in list_brains:
if brain["id"] not in seen_brain_ids:
unique_list_brains.append(brain)
seen_brain_ids.add(brain["id"])
for brain in list_brains:
if brain["id"] not in seen_brain_ids:
unique_list_brains.append(brain)
seen_brain_ids.add(brain["id"])
metadata["close_brains"] = unique_list_brains[:5]
metadata["close_brains"] = unique_list_brains[:5]
if list_brains and not brain_to_use:
brain_id_to_use = list_brains[0]["id"]
brain_to_use = self.get_brain_by_id(brain_id_to_use)
if list_brains and not brain_to_use:
brain_id_to_use = list_brains[0]["id"]
brain_to_use = self.get_brain_by_id(brain_id_to_use)
return brain_to_use, metadata
@ -323,7 +322,9 @@ class BrainService:
def update_brain_last_update_time(self, brain_id: UUID):
self.brain_repository.update_brain_last_update_time(brain_id)
def get_brain_details(self, brain_id: UUID, user_id: UUID = None) -> BrainEntity | None:
def get_brain_details(
self, brain_id: UUID, user_id: UUID = None
) -> BrainEntity | None:
brain = self.brain_repository.get_brain_details(brain_id)
if brain == None:
return None
@ -334,14 +335,14 @@ class BrainService:
brain_id, user_id
)
)
if (brain.integration):
if brain.integration:
brain.integration_description = (
self.integration_description_repository.get_integration_description(
brain.integration.integration_id
)
)
return brain
def get_connected_brains(self, brain_id: UUID) -> list[BrainEntity]:

6
backend/modules/chat/controller/chat_routes.py
View File

@ -61,7 +61,6 @@ def get_answer_generator(
current_user: UserIdentity,
):
chat_instance = BrainfulChat()
chat_instance.validate_authorization(user_id=current_user.id, brain_id=brain_id)
user_usage = UserUsage(
id=current_user.id,
@ -247,11 +246,6 @@ async def create_stream_question_handler(
chat_instance = BrainfulChat()
chat_instance.validate_authorization(user_id=current_user.id, brain_id=brain_id)
user_usage = UserUsage(
id=current_user.id,
email=current_user.email,
)
logger.info(
f"Creating question for chat {chat_id} with brain {brain_id} of type {type(brain_id)}"
)