quivr/backend/modules/brain/service/brain_service.py
Zineb El Bachiri 436e49a5e7
refactor: chat for multibrains (#1812)
# Description

- Chat Module
- External Api Secrets Interface, exposed through brain service
2023-12-04 18:38:54 +01:00

236 lines
8.4 KiB
Python

from typing import Optional
from uuid import UUID
from fastapi import HTTPException
from modules.brain.dto.inputs import BrainUpdatableProperties, CreateBrainProperties
from modules.brain.entity.brain_entity import BrainEntity, BrainType, PublicBrain
from modules.brain.repository.brains import Brains
from modules.brain.repository.brains_users import BrainsUsers
from modules.brain.repository.brains_vectors import BrainsVectors
from modules.brain.repository.external_api_secrets import ExternalApiSecrets
from modules.brain.repository.interfaces.brains_interface import BrainsInterface
from modules.brain.repository.interfaces.brains_users_interface import (
BrainsUsersInterface,
)
from modules.brain.repository.interfaces.brains_vectors_interface import (
BrainsVectorsInterface,
)
from modules.brain.repository.interfaces.external_api_secrets_interface import (
ExternalApiSecretsInterface,
)
from modules.knowledge.service.knowledge_service import KnowledgeService
from repository.api_brain_definition.add_api_brain_definition import (
add_api_brain_definition,
)
from repository.api_brain_definition.delete_api_brain_definition import (
delete_api_brain_definition,
)
from repository.api_brain_definition.get_api_brain_definition import (
get_api_brain_definition,
)
from repository.api_brain_definition.update_api_brain_definition import (
update_api_brain_definition,
)
knowledge_service = KnowledgeService()
class BrainService:
brain_repository: BrainsInterface
brain_user_repository: BrainsUsersInterface
brain_vector_repository: BrainsVectorsInterface
external_api_secrets_repository: ExternalApiSecretsInterface
def __init__(self):
self.brain_repository = Brains()
self.brain_user_repository = BrainsUsers()
self.brain_vector = BrainsVectors()
self.external_api_secrets_repository = ExternalApiSecrets()
def get_brain_by_id(self, brain_id: UUID):
return self.brain_repository.get_brain_by_id(brain_id)
def create_brain(
self,
user_id: UUID,
brain: Optional[CreateBrainProperties],
) -> BrainEntity:
if brain == None:
brain = CreateBrainProperties() # type: ignore model and brain_definition
if brain.brain_type == BrainType.API:
if brain.brain_definition is None:
raise HTTPException(
status_code=404, detail="Brain definition not found"
)
if brain.brain_definition.url is None:
raise HTTPException(status_code=404, detail="Brain url not found")
if brain.brain_definition.method is None:
raise HTTPException(status_code=404, detail="Brain method not found")
created_brain = self.brain_repository.create_brain(brain)
if brain.brain_type == BrainType.API and brain.brain_definition is not None:
add_api_brain_definition(
brain_id=created_brain.brain_id,
api_brain_definition=brain.brain_definition,
)
secrets_values = brain.brain_secrets_values
for secret_name in secrets_values:
self.external_api_secrets_repository.create_secret(
user_id=user_id,
brain_id=created_brain.brain_id,
secret_name=secret_name,
secret_value=secrets_values[secret_name],
)
return created_brain
def delete_brain_secrets_values(self, brain_id: UUID) -> None:
brain_definition = get_api_brain_definition(brain_id=brain_id)
if brain_definition is None:
raise HTTPException(status_code=404, detail="Brain definition not found.")
secrets = brain_definition.secrets
if len(secrets) > 0:
brain_users = self.brain_user_repository.get_brain_users(brain_id=brain_id)
for user in brain_users:
for secret in secrets:
self.external_api_secrets_repository.delete_secret(
user_id=user.user_id,
brain_id=brain_id,
secret_name=secret.name,
)
def delete_brain(self, brain_id: UUID) -> dict[str, str]:
brain_to_delete = self.get_brain_by_id(brain_id=brain_id)
if brain_to_delete is None:
raise HTTPException(status_code=404, detail="Brain not found.")
if brain_to_delete.brain_type == BrainType.API:
self.delete_brain_secrets_values(
brain_id=brain_id,
)
delete_api_brain_definition(brain_id=brain_id)
else:
knowledge_service.remove_brain_all_knowledge(brain_id)
self.brain_vector.delete_brain_vector(str(brain_id))
self.brain_user_repository.delete_brain_users(str(brain_id))
self.brain_repository.delete_brain(str(brain_id)) # type: ignore
return {"message": "Brain deleted."}
def get_brain_prompt_id(self, brain_id: UUID) -> UUID | None:
brain = self.get_brain_by_id(brain_id)
prompt_id = brain.prompt_id if brain else None
return prompt_id
def update_brain_by_id(
self, brain_id: UUID, brain_new_values: BrainUpdatableProperties
) -> BrainEntity:
"""Update a prompt by id"""
existing_brain = self.brain_repository.get_brain_by_id(brain_id)
if existing_brain is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
brain_update_answer = self.brain_repository.update_brain_by_id(
brain_id,
brain=BrainUpdatableProperties(
**brain_new_values.dict(exclude={"brain_definition"})
),
)
if brain_update_answer is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
if (
brain_update_answer.brain_type == BrainType.API
and brain_new_values.brain_definition
):
existing_brain_secrets_definition = (
existing_brain.brain_definition.secrets
if existing_brain.brain_definition
else None
)
brain_new_values_secrets_definition = (
brain_new_values.brain_definition.secrets
if brain_new_values.brain_definition
else None
)
should_remove_existing_secrets_values = (
existing_brain_secrets_definition
and brain_new_values_secrets_definition
and existing_brain_secrets_definition
!= brain_new_values_secrets_definition
)
if should_remove_existing_secrets_values:
self.delete_brain_secrets_values(brain_id=brain_id)
update_api_brain_definition(
brain_id,
api_brain_definition=brain_new_values.brain_definition,
)
if brain_update_answer is None:
raise HTTPException(
status_code=404,
detail=f"Brain with id {brain_id} not found",
)
self.brain_repository.update_brain_last_update_time(brain_id)
return brain_update_answer
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) -> BrainEntity | None:
brain = self.brain_repository.get_brain_details(brain_id)
# id ?
if brain == None:
return None
if brain.brain_type == BrainType.API:
brain_definition = get_api_brain_definition(brain_id)
brain.brain_definition = brain_definition
return brain
def get_public_brains(self) -> list[PublicBrain]:
return self.brain_repository.get_public_brains()
def update_secret_value(
self,
user_id: UUID,
brain_id: UUID,
secret_name: str,
secret_value: str,
) -> None:
"""Update an existing secret."""
self.external_api_secrets_repository.delete_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret_name,
)
self.external_api_secrets_repository.create_secret(
user_id=user_id,
brain_id=brain_id,
secret_name=secret_name,
secret_value=secret_value,
)