codekingpro/portable-devtools

1"""Models for the PebbloRetrievalQA chain."""2 3from typing import Any, List, Optional, Union4 5from pydantic import BaseModel6 7 8class AuthContext(BaseModel):9    """Class for an authorization context."""10 11    name: Optional[str] = None12    user_id: str13    user_auth: List[str]14    """List of user authorizations, which may include their User ID and 15    the groups they are part of"""16 17 18class SemanticEntities(BaseModel):19    """Class for a semantic entity filter."""20 21    deny: List[str]22 23 24class SemanticTopics(BaseModel):25    """Class for a semantic topic filter."""26 27    deny: List[str]28 29 30class SemanticContext(BaseModel):31    """Class for a semantic context."""32 33    pebblo_semantic_entities: Optional[SemanticEntities] = None34    pebblo_semantic_topics: Optional[SemanticTopics] = None35 36    def __init__(self, **data: Any) -> None:37        super().__init__(**data)38 39        # Validate semantic_context40        if (41            self.pebblo_semantic_entities is None42            and self.pebblo_semantic_topics is None43        ):44            raise ValueError(45                "semantic_context must contain 'pebblo_semantic_entities' or "46                "'pebblo_semantic_topics'"47            )48 49 50class ChainInput(BaseModel):51    """Input for PebbloRetrievalQA chain."""52 53    query: str54    auth_context: Optional[AuthContext] = None55    semantic_context: Optional[SemanticContext] = None56 57    def dict(self, **kwargs: Any) -> dict:58        base_dict = super().dict(**kwargs)59        # Keep auth_context and semantic_context as it is(Pydantic models)60        base_dict["auth_context"] = self.auth_context61        base_dict["semantic_context"] = self.semantic_context62        return base_dict63 64 65class Runtime(BaseModel):66    """67    OS, language details68    """69 70    type: Optional[str] = ""71    host: str72    path: str73    ip: Optional[str] = ""74    platform: str75    os: str76    os_version: str77    language: str78    language_version: str79    runtime: Optional[str] = ""80 81 82class Framework(BaseModel):83    """84    Langchain framework details85    """86 87    name: str88    version: str89 90 91class Model(BaseModel):92    vendor: Optional[str]93    name: Optional[str]94 95 96class PkgInfo(BaseModel):97    project_home_page: Optional[str]98    documentation_url: Optional[str]99    pypi_url: Optional[str]100    liscence_type: Optional[str]101    installed_via: Optional[str]102    location: Optional[str]103 104 105class VectorDB(BaseModel):106    name: Optional[str] = None107    version: Optional[str] = None108    location: Optional[str] = None109    embedding_model: Optional[str] = None110 111 112class ChainInfo(BaseModel):113    name: str114    model: Optional[Model]115    vector_dbs: Optional[List[VectorDB]]116 117 118class App(BaseModel):119    name: str120    owner: str121    description: Optional[str]122    runtime: Runtime123    framework: Framework124    chains: List[ChainInfo]125    plugin_version: str126    client_version: Framework127 128 129class Context(BaseModel):130    retrieved_from: Optional[str]131    doc: Optional[str]132    vector_db: str133    pb_checksum: Optional[str]134 135 136class Prompt(BaseModel):137    data: Optional[Union[list, str]]138    entityCount: Optional[int] = None139    entities: Optional[dict] = None140    prompt_gov_enabled: Optional[bool] = None141 142 143class Qa(BaseModel):144    name: str145    context: Union[List[Optional[Context]], Optional[Context]]146    prompt: Optional[Prompt]147    response: Optional[Prompt]148    prompt_time: str149    user: str150    user_identities: Optional[List[str]]151    classifier_location: str152