mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-24 07:36:39 +03:00
436e49a5e7
# Description - Chat Module - External Api Secrets Interface, exposed through brain service
236 lines
8.4 KiB
Python
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,
|
|
)
|