mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-12-18 11:51:41 +03:00
fb3aae27f1
# Description Please include a summary of the changes and the related issue. Please also include relevant motivation and context. ## Checklist before requesting a review Please delete options that are not relevant. - [ ] My code follows the style guidelines of this project - [ ] I have performed a self-review of my code - [ ] I have commented hard-to-understand areas - [ ] I have ideally added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [ ] Any dependent changes have been merged ## Screenshots (if appropriate):
83 lines
1.9 KiB
Python
83 lines
1.9 KiB
Python
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import List, Optional
|
|
from uuid import UUID
|
|
|
|
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinitionEntity
|
|
from modules.brain.entity.integration_brain import (
|
|
IntegrationDescriptionEntity,
|
|
IntegrationEntity,
|
|
)
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class BrainType(str, Enum):
|
|
DOC = "doc"
|
|
API = "api"
|
|
COMPOSITE = "composite"
|
|
INTEGRATION = "integration"
|
|
|
|
|
|
class BrainEntity(BaseModel):
|
|
brain_id: UUID
|
|
name: str
|
|
description: Optional[str] = None
|
|
temperature: Optional[float] = None
|
|
model: Optional[str] = None
|
|
max_tokens: Optional[int] = None
|
|
status: Optional[str] = None
|
|
prompt_id: Optional[UUID] = None
|
|
last_update: datetime
|
|
brain_type: BrainType
|
|
brain_definition: Optional[ApiBrainDefinitionEntity] = None
|
|
connected_brains_ids: Optional[List[UUID]] = None
|
|
raw: Optional[bool] = None
|
|
jq_instructions: Optional[str] = None
|
|
integration: Optional[IntegrationEntity] = None
|
|
integration_description: Optional[IntegrationDescriptionEntity] = None
|
|
|
|
@property
|
|
def id(self) -> UUID:
|
|
return self.brain_id
|
|
|
|
def dict(self, **kwargs):
|
|
data = super().dict(
|
|
**kwargs,
|
|
)
|
|
data["id"] = self.id
|
|
return data
|
|
|
|
|
|
class PublicBrain(BaseModel):
|
|
id: UUID
|
|
name: str
|
|
description: Optional[str] = None
|
|
number_of_subscribers: int = 0
|
|
last_update: str
|
|
brain_type: BrainType
|
|
brain_definition: Optional[ApiBrainDefinitionEntity] = None
|
|
|
|
|
|
class RoleEnum(str, Enum):
|
|
Viewer = "Viewer"
|
|
Editor = "Editor"
|
|
Owner = "Owner"
|
|
|
|
|
|
class BrainUser(BaseModel):
|
|
id: UUID
|
|
user_id: UUID
|
|
rights: RoleEnum
|
|
default_brain: bool = False
|
|
|
|
|
|
class MinimalUserBrainEntity(BaseModel):
|
|
id: UUID
|
|
name: str
|
|
rights: RoleEnum
|
|
status: str
|
|
brain_type: BrainType
|
|
description: str
|
|
integration_logo_url: str
|
|
max_files: int
|