Source code for mb.entities.outputs

from datetime import datetime
from typing import Annotated, Literal, Self

from pydantic import BaseModel, Field, JsonValue, model_validator




class Environment(BaseModel, extra="allow"):
    user_time: datetime | None = None
    """The time the user sent the request."""





class ToolCall(BaseModel):
    event: Literal["tool_call"] = "tool_call"
    """The type of event."""

    id: str
    """The ID of the tool call."""

    tool: str
    """The name of the tool called. Must correspond to one of the tools in the benchmark description file."""

    params: dict[str, JsonValue]
    """The parameters passed to the tool."""





class ToolResult(BaseModel):
    event: Literal["tool_result"] = "tool_result"
    """The type of event."""

    id: str
    """The ID of the tool call that yielded this result."""

    result: JsonValue
    """The result of the tool call."""



class RetrievalResult(BaseModel, extra="allow"):
    id: str
    """The application-specific ID of the chunk."""

    page_content: str
    """The content of the chunk."""


class RetrievalResults(BaseModel, extra="allow"):
    event: Literal["retriever"] = "retriever"
    """The type of event."""

    outputs: list[RetrievalResult]
    """The results of the retrieval operation."""


TraceEvent = Annotated[ToolCall | ToolResult | RetrievalResults, Field(discriminator="event")]


class Citation(BaseModel):
    document_id: str
    """The ID of the document that was cited. Has to be unique within the retrieval results of the trace."""

    span_from: int
    """The start index of the span that was cited, inclusive."""

    span_to: int
    """The end index of the span that was cited, exclusive."""




class Outputs(BaseModel, extra="allow"):
    response: str
    """The agent's response to the user."""

    citations: list[Citation] | None = None
    """The citations the agent made to justify its response."""

    environment: Environment | None = None
    """Additional environment information, such as the time the user sent the request."""

    trace: list[TraceEvent] = []
    """The trace of the agent's execution events, such as tool calls, tool call results, search, etc.

    This is required for some assertions to work correctly, such as tool-call correctness assertions.
    """

    @model_validator(mode="after")
    def _validate_citations(self) -> Self:
        if not self.citations:
            return self

        retrieval_ids = {
            result.id for event in self.trace if isinstance(event, RetrievalResults) for result in event.outputs
        }

        missing = {citation.document_id for citation in self.citations} - retrieval_ids
        if missing:
            raise ValueError(f"Citations reference unknown document ids: {', '.join(missing)}.")

        return self