quivr/backend/modules/brain/entity/brain_entity.py
Stan Girard 25a53a254b
feat: 🎸 brains (#2048)
update vectors at launch if missign

# 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):
2024-01-21 17:03:25 -08:00

71 lines
1.4 KiB
Python

from enum import Enum
from typing import List, Optional
from uuid import UUID
from modules.brain.entity.api_brain_definition_entity import ApiBrainDefinitionEntity
from pydantic import BaseModel
class BrainType(str, Enum):
DOC = "doc"
API = "api"
COMPOSITE = "composite"
class BrainEntity(BaseModel):
brain_id: UUID
name: str
description: Optional[str]
temperature: Optional[float]
model: Optional[str]
max_tokens: Optional[int]
status: Optional[str]
prompt_id: Optional[UUID]
last_update: str
brain_type: BrainType
brain_definition: Optional[ApiBrainDefinitionEntity]
connected_brains_ids: Optional[List[UUID]]
@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]
number_of_subscribers: int = 0
last_update: str
brain_type: BrainType
brain_definition: Optional[ApiBrainDefinitionEntity]
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