事件类型

class BaseEvent(BaseModel):
    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        # copy_on_model_validation = "deep"  # not supported in Pydantic V2...
    )
    timestamp: datetime = Field(default_factory=lambda: datetime.now())
    id_: str = Field(default_factory=lambda: str(uuid4()))
    span_id: Optional[str] = Field(default_factory=active_span_id.get)  # type: ignore
    tags: Dict[str, Any] = Field(default={})

    @classmethod
    def class_name(cls) -> str:
        """Return class name."""
        return "BaseEvent"

    def dict(self, **kwargs: Any) -> Dict[str, Any]:
        """Keep for backwards compatibility."""
        return self.model_dump(**kwargs)

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
        data = super().model_dump(**kwargs)
        data["class_name"] = self.class_name()
        return data

class AgentChatWithStepEndEvent(BaseEvent):
    """
    AgentChatWithStepEndEvent.

    Args:
        response (Optional[AGENT_CHAT_RESPONSE_TYPE]): Agent chat response.

    """

    response: Optional[AGENT_CHAT_RESPONSE_TYPE]

    @model_validator(mode="before")
    @classmethod
    def validate_response(cls: Any, values: Any) -> Any:
        """Validate response."""
        response = values.get("response")
        if response is None:
            pass
        elif not isinstance(response, AgentChatResponse) and not isinstance(
            response, StreamingAgentChatResponse
        ):
            raise ValueError(
                "response must be of type AgentChatResponse or StreamingAgentChatResponse"
            )

        return values

    @field_validator("response", mode="before")
    @classmethod
    def validate_response_type(cls: Any, response: Any) -> Any:
        """Validate response type."""
        if response is None:
            return response
        if not isinstance(response, AgentChatResponse) and not isinstance(
            response, StreamingAgentChatResponse
        ):
            raise ValueError(
                "response must be of type AgentChatResponse or StreamingAgentChatResponse"
            )
        return response

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "AgentChatWithStepEndEvent"

class AgentChatWithStepStartEvent(BaseEvent):
    """
    AgentChatWithStepStartEvent.

    Args:
        user_msg (str): User input message.

    """

    user_msg: str

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "AgentChatWithStepStartEvent"

class AgentRunStepEndEvent(BaseEvent):
    """
    AgentRunStepEndEvent.

    Args:
        step_output (Any): Task step output.

    """

    step_output: Any

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "AgentRunStepEndEvent"

class AgentRunStepStartEvent(BaseEvent):
    """
    AgentRunStepStartEvent.

    Args:
        task_id (str): Task ID.
        step (Optional[Any]): Task step.
        input (Optional[str]): Optional input.

    """

    task_id: str
    step: Optional[Any]
    input: Optional[str]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "AgentRunStepStartEvent"

class AgentToolCallEvent(BaseEvent):
    """
    AgentToolCallEvent.

    Args:
        arguments (str): Arguments.
        tool (ToolMetadata): Tool metadata.

    """

    arguments: str
    tool: ToolMetadata

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "AgentToolCallEvent"

class StreamChatDeltaReceivedEvent(BaseEvent):
    """
    StreamChatDeltaReceivedEvent.

    Args:
        delta (str): Delta received from the stream chat.

    """

    delta: str

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "StreamChatDeltaReceivedEvent"

class StreamChatEndEvent(BaseEvent):
    """
    StreamChatEndEvent.

    Fired at the end of writing to the stream chat-engine queue.
    """

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "StreamChatEndEvent"

class StreamChatErrorEvent(BaseEvent):
    """
    StreamChatErrorEvent.

    Fired when an exception is raised during the stream chat-engine operation.

    Args:
        exception (Exception): Exception raised during the stream chat operation.

    """

    exception: Exception

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "StreamChatErrorEvent"

class StreamChatStartEvent(BaseEvent):
    """
    StreamChatStartEvent.

    Fired at the start of writing to the stream chat-engine queue.
    """

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "StreamChatStartEvent"

class EmbeddingEndEvent(BaseEvent):
    """
    EmbeddingEndEvent.

    Args:
        chunks (List[str]): List of chunks.
        embeddings (List[List[float]]): List of embeddings.

    """

    chunks: List[str]
    embeddings: List[List[float]]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "EmbeddingEndEvent"

class EmbeddingStartEvent(BaseEvent):
    """
    EmbeddingStartEvent.

    Args:
        model_dict (dict): Model dictionary containing details about the embedding model.

    """

    model_config = ConfigDict(protected_namespaces=("pydantic_model_",))
    model_dict: dict

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "EmbeddingStartEvent"

class LLMChatEndEvent(BaseEvent):
    """
    LLMChatEndEvent.

    Args:
        messages (List[ChatMessage]): List of chat messages.
        response (Optional[ChatResponse]): Last chat response.

    """

    messages: List[ChatMessage]
    response: Optional[ChatResponse]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMChatEndEvent"

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
        if self.response is not None and isinstance(self.response.raw, BaseModel):
            self.response.raw = self.response.raw.model_dump()

        return super().model_dump(**kwargs)

class LLMChatStartEvent(BaseEvent):
    """
    LLMChatStartEvent.

    Args:
        messages (List[ChatMessage]): List of chat messages.
        additional_kwargs (dict): Additional keyword arguments.
        model_dict (dict): Model dictionary.

    """

    model_config = ConfigDict(protected_namespaces=("pydantic_model_",))
    messages: List[ChatMessage]
    additional_kwargs: dict
    model_dict: dict

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMChatStartEvent"

class LLMCompletionEndEvent(BaseEvent):
    """
    LLMCompletionEndEvent.

    Args:
        prompt (str): The prompt to be completed.
        response (CompletionResponse): Completion response.

    """

    prompt: str
    response: CompletionResponse

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMCompletionEndEvent"

    def model_dump(self, **kwargs: Any) -> Dict[str, Any]:
        if isinstance(self.response.raw, BaseModel):
            self.response.raw = self.response.raw.model_dump()

        return super().model_dump(**kwargs)

class LLMCompletionStartEvent(BaseEvent):
    """
    LLMCompletionStartEvent.

    Args:
        prompt (str): The prompt to be completed.
        additional_kwargs (dict): Additional keyword arguments.
        model_dict (dict): Model dictionary.

    """

    model_config = ConfigDict(protected_namespaces=("pydantic_model_",))
    prompt: str
    additional_kwargs: dict
    model_dict: dict

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMCompletionStartEvent"

class LLMPredictEndEvent(BaseEvent):
    """
    LLMPredictEndEvent.

    The result of an llm.predict() call.

    Args:
        output (str): Output.

    """

    output: str

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMPredictEndEvent"

class LLMPredictStartEvent(BaseEvent):
    """
    LLMPredictStartEvent.

    Args:
        template (BasePromptTemplate): Prompt template.
        template_args (Optional[dict]): Prompt template arguments.

    """

    template: SerializeAsAny[BasePromptTemplate]
    template_args: Optional[dict]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "LLMPredictStartEvent"

class QueryEndEvent(BaseEvent):
    """
    QueryEndEvent.

    Args:
        query (QueryType): Query as a string or query bundle.
        response (RESPONSE_TYPE): Response.

    """

    query: QueryType
    response: RESPONSE_TYPE

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "QueryEndEvent"

class QueryStartEvent(BaseEvent):
    """
    QueryStartEvent.

    Args:
        query (QueryType): Query as a string or query bundle.

    """

    query: QueryType

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "QueryStartEvent"

class RetrievalEndEvent(BaseEvent):
    """
    RetrievalEndEvent.

    Args:
        str_or_query_bundle (QueryType): Query bundle.
        nodes (List[NodeWithScore]): List of nodes with scores.

    """

    str_or_query_bundle: QueryType
    nodes: List[NodeWithScore]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "RetrievalEndEvent"

class RetrievalStartEvent(BaseEvent):
    """
    RetrievalStartEvent.

    Args:
        str_or_query_bundle (QueryType): Query bundle.

    """

    str_or_query_bundle: QueryType

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "RetrievalStartEvent"

class GetResponseEndEvent(BaseEvent):
    """GetResponseEndEvent."""

    # TODO: consumes the first chunk of generators??
    # response: RESPONSE_TEXT_TYPE

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "GetResponseEndEvent"

class GetResponseStartEvent(BaseEvent):
    """
    GetResponseStartEvent.

    Args:
        query_str (str): Query string.
        text_chunks (List[str]): List of text chunks.

    """

    query_str: str
    text_chunks: List[str]

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "GetResponseStartEvent"

class SynthesizeEndEvent(BaseEvent):
    """
    SynthesizeEndEvent.

    Args:
        query (QueryType): Query as a string or query bundle.
        response (RESPONSE_TYPE): Response.

    """

    query: QueryType
    response: RESPONSE_TYPE

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "SynthesizeEndEvent"

class SynthesizeStartEvent(BaseEvent):
    """
    SynthesizeStartEvent.

    Args:
        query (QueryType): Query as a string or query bundle.

    """

    query: QueryType

    @classmethod
    def class_name(cls) -> str:
        """Class name."""
        return "SynthesizeStartEvent"