Source code for langchain_community.chains.pebblo_retrieval.models

"""PebbloRetrievalQA链的模型。"""

from typing import Any, List, Optional

from langchain_core.pydantic_v1 import BaseModel


class AuthContext(BaseModel):
    """用于授权上下文的类。"""

    name: Optional[str] = None
    user_id: str
    user_auth: List[str]
    """用户授权列表，可能包括他们的用户ID和所属的用户组。"""


class SemanticEntities(BaseModel):
    """用于语义实体过滤的类。"""

    deny: List[str]


class SemanticTopics(BaseModel):
    """用于语义主题过滤的类。"""

    deny: List[str]


class SemanticContext(BaseModel):
    """用于语义上下文的类。"""

    pebblo_semantic_entities: Optional[SemanticEntities] = None
    pebblo_semantic_topics: Optional[SemanticTopics] = None

    def __init__(self, **data: Any) -> None:
        super().__init__(**data)

        # Validate semantic_context
        if (
            self.pebblo_semantic_entities is None
            and self.pebblo_semantic_topics is None
        ):
            raise ValueError(
                "semantic_context must contain 'pebblo_semantic_entities' or "
                "'pebblo_semantic_topics'"
            )


class ChainInput(BaseModel):
    """PebbloRetrievalQA链的输入。"""

    query: str
    auth_context: Optional[AuthContext] = None
    semantic_context: Optional[SemanticContext] = None

    def dict(self, **kwargs: Any) -> dict:
        base_dict = super().dict(**kwargs)
        # Keep auth_context and semantic_context as it is(Pydantic models)
        base_dict["auth_context"] = self.auth_context
        base_dict["semantic_context"] = self.semantic_context
        return base_dict