rasa-pro 3.14.0.dev5__py3-none-any.whl → 3.14.0.dev7__py3-none-any.whl

rasa/builder/copilot/models.py

@@ -1,431 +0,0 @@
-import json
-from abc import ABC, abstractmethod
-from enum import Enum
-from typing import Any, Dict, List, Literal, Optional, Union
-
-import structlog
-from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
-from pydantic import BaseModel, Field, model_validator
-from typing_extensions import Annotated
-
-from rasa.builder.copilot.constants import (
-    ROLE_ASSISTANT,
-    ROLE_COPILOT,
-    ROLE_COPILOT_INTERNAL,
-    ROLE_USER,
-)
-from rasa.builder.models import ServerSentEvent
-from rasa.builder.shared.tracker_context import TrackerContext
-
-structlogger = structlog.get_logger()
-
-
-class ResponseCompleteness(Enum):
-    """Enum for response completeness levels."""
-
-    TOKEN = "token"  # Streaming token/chunk
-    COMPLETE = "complete"  # Complete response (e.g., templated responses)
-
-
-class ResponseCategory(Enum):
-    """Enum for different categories of responses."""
-
-    # Copilot generated content
-    COPILOT = "copilot"
-    REFERENCE = "reference"
-    REFERENCE_ENTRY = "reference_entry"
-    # When Copilot detects a roleplay request / intent
-    ROLEPLAY_DETECTION = "roleplay_detection"
-    # When Copilot detects an out-of-scope request
-    OUT_OF_SCOPE_DETECTION = "out_of_scope_detection"
-    # When Copilot does not understand what caused the error
-    ERROR_FALLBACK = "error_fallback"
-    # When a policy violation is detected
-    GUARDRAILS_POLICY_VIOLATION = "guardrails_policy_violation"
-    # When Copilot analyzes error logs and provides suggestions
-    TRAINING_ERROR_LOG_ANALYSIS = "training_error_log_analysis"
-    E2E_TESTING_ERROR_LOG_ANALYSIS = "e2e_testing_error_log_analysis"
-
-
-class BaseContent(BaseModel):
-    type: str
-
-
-class LinkContent(BaseContent):
-    type: Literal["link"]
-    url: str
-    label: str
-
-
-class ButtonContent(BaseContent):
-    type: Literal["button"]
-    payload: str
-    label: str
-
-
-class TextContent(BaseContent):
-    type: Literal["text"]
-    text: str
-
-
-class CodeContent(BaseContent):
-    type: Literal["code"]
-    text: str
-
-
-class FileContent(BaseContent):
-    type: Literal["file"]
-    file_path: str
-    file_content: str
-
-
-class LogContent(BaseContent):
-    type: Literal["log"]
-    content: str = Field(..., description="Logs, error messages, stack traces, etc.")
-    context: Optional[str] = Field(
-        None,
-        description=(
-            "Additional, optional context description for the logs "
-            "(e.g., 'training session', 'e2e testing run', 'deployment process')"
-        ),
-    )
-    metadata: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Additional metadata about the log.",
-    )
-
-
-class EventContent(BaseModel):
-    type: Literal["event"]
-    event: str = Field(..., description="The event's type_name")
-
-    event_data: Dict[str, Any] = Field(
-        default_factory=dict, description="Contains event-specific data fields."
-    )
-
-    @model_validator(mode="before")
-    @classmethod
-    def _collect_event_data(cls, data: Dict[str, Any]) -> Dict[str, Any]:
-        generic = {"type", "event"}
-        data["event_data"] = {
-            key: data.pop(key) for key in list(data.keys()) if key not in generic
-        }
-        return data
-
-    class Config:
-        extra = "forbid"
-
-
-ContentBlock = Annotated[
-    Union[
-        TextContent,
-        CodeContent,
-        FileContent,
-        LogContent,
-        EventContent,
-        LinkContent,
-        ButtonContent,
-    ],
-    Field(
-        discriminator="type",
-        description=(
-            "The content of the message. "
-            "The content is expected to be a list of content blocks. "
-            "The content blocks are expected to be one of the following types: "
-            "text, link, code, or file."
-        ),
-    ),
-]
-
-
-class CopilotChatMessage(BaseModel):
-    """Model for a single chat messages between the user and the copilot."""
-
-    role: str = Field(
-        ...,
-        pattern=f"^({ROLE_USER}|{ROLE_COPILOT}|{ROLE_COPILOT_INTERNAL})$",
-        description="The role of the message sender.",
-    )
-    content: List[ContentBlock]
-    response_category: Optional[ResponseCategory] = Field(
-        None,
-        description=(
-            "The category/source of this message. For user role messages, only `None` "
-            "or `GUARDRAILS_POLICY_VIOLATION` are allowed. For copilot role messages, "
-            "any category is permitted."
-        ),
-    )
-
-    @model_validator(mode="after")
-    def validate_response_category_for_role(self) -> "CopilotChatMessage":
-        """Validate value of response_category for the role of the message.
-
-        For 'user' role messages, only None or GUARDRAILS_POLICY_VIOLATION are allowed.
-        For 'copilot' role messages, any category is permitted.
-        For 'rasa_internal' role messages, any category is permitted.
-        """
-        if (
-            self.role == ROLE_USER
-            and self.response_category is not None
-            and self.response_category != ResponseCategory.GUARDRAILS_POLICY_VIOLATION
-        ):
-            message = (
-                f"User role messages can only have response_category of `None` or "
-                f"`{ResponseCategory.GUARDRAILS_POLICY_VIOLATION}`, "
-                f"got `{self.response_category}`."
-            )
-            structlogger.error(
-                "copilot_chat_message.validate_response_category_for_role"
-                ".invalid_response_category",
-                event_info=message,
-                response_category=self.response_category,
-                role=self.role,
-            )
-            raise ValueError(message)
-
-        return self
-
-    def get_text_content(self) -> str:
-        """Concatenate all 'text' content blocks into a single string."""
-        return "\n".join(
-            content_block.text
-            for content_block in self.content
-            if isinstance(content_block, TextContent)
-        )
-
-    def get_log_content(self) -> str:
-        """Concatenate all 'log' content blocks into a single string."""
-        return "\n".join(
-            content_block.content
-            for content_block in self.content
-            if isinstance(content_block, LogContent)
-        )
-
-    def to_openai_format(self) -> Dict[str, Any]:
-        """Convert to OpenAI message format for API calls."""
-        role_to_openai_format = {
-            ROLE_USER: self._user_message_to_openai_format,
-            ROLE_COPILOT: self._copilot_message_to_openai_format,
-            ROLE_COPILOT_INTERNAL: self._copilot_message_to_openai_format,
-        }
-        return role_to_openai_format[self.role]()
-
-    def _user_message_to_openai_format(self) -> Dict[str, Any]:
-        role = self._map_role_to_openai()
-        content = self.get_text_content()
-        return {"role": role, "content": content}
-
-    def _copilot_message_to_openai_format(self) -> Dict[str, Any]:
-        role = self._map_role_to_openai()
-        content_blocks = [block.model_dump() for block in self.content]
-        content = json.dumps(content_blocks)
-        return {"role": role, "content": content}
-
-    def _map_role_to_openai(self) -> str:
-        """Map internal roles to OpenAI-compatible roles."""
-        role_mapping = {
-            ROLE_USER: ROLE_USER,
-            ROLE_COPILOT: ROLE_ASSISTANT,
-            ROLE_COPILOT_INTERNAL: ROLE_USER,
-        }
-        if self.role not in role_mapping.keys():
-            structlogger.error(
-                "copilot_chat_message.to_openai_format.invalid_role",
-                event_info=(
-                    f"Invalid role: `{self.role}`. "
-                    f"Only {', '.join(role_mapping.keys())} roles are supported."
-                ),
-                role=self.role,
-            )
-            raise ValueError(f"Invalid role: {self.role}")
-
-        return role_mapping[self.role]
-
-
-class CopilotRequest(BaseModel):
-    """Request model for the copilot endpoint."""
-
-    copilot_chat_history: List[CopilotChatMessage] = Field(
-        ...,
-        description=(
-            "The chat history between the user and the copilot. "
-            "Used to generate a new response based on the previous conversation."
-        ),
-    )
-    session_id: str = Field(
-        ...,
-        description=(
-            "The session ID of chat session with the assistant. "
-            "Used to fetch the conversation from the tracker."
-        ),
-    )
-
-    @property
-    def last_message(self) -> Optional[CopilotChatMessage]:
-        """Get the last message from the copilot chat history."""
-        if not self.copilot_chat_history:
-            return None
-        return self.copilot_chat_history[-1]
-
-
-class CopilotOutput(BaseModel, ABC):
-    """Base class for response events."""
-
-    response_completeness: ResponseCompleteness = Field(
-        description=(
-            "Indicates whether this is a streaming token (TOKEN) or a complete "
-            "response (COMPLETE)"
-        ),
-    )
-    response_category: ResponseCategory = Field(
-        description=(
-            "The category/source of this response. Each response type has a fixed "
-            "category that cannot be changed. "
-        ),
-        frozen=True,
-    )
-
-    @abstractmethod
-    def to_sse_event(self) -> ServerSentEvent:
-        """Convert to SSE event format."""
-        pass
-
-
-class GeneratedContent(CopilotOutput):
-    """Represents generated content from the LLM to be streamed."""
-
-    content: str
-    response_category: ResponseCategory = Field(frozen=True)
-    response_completeness: ResponseCompleteness = ResponseCompleteness.TOKEN
-
-    def to_sse_event(self) -> ServerSentEvent:
-        """Convert to SSE event format."""
-        return ServerSentEvent(
-            event="copilot_response",
-            data={
-                "content": self.content,
-                "response_category": self.response_category.value,
-                "completeness": self.response_completeness.value,
-            },
-        )
-
-
-class ReferenceEntry(CopilotOutput):
-    """Represents a reference entry with title and url."""
-
-    index: int
-    title: str
-    url: str
-    response_category: ResponseCategory = Field(
-        default=ResponseCategory.REFERENCE_ENTRY,
-        frozen=True,
-    )
-    response_completeness: ResponseCompleteness = ResponseCompleteness.COMPLETE
-
-    @model_validator(mode="after")
-    def validate_response_category(self) -> "ReferenceEntry":
-        """Validate that response_category has the correct default value."""
-        if self.response_category != ResponseCategory.REFERENCE_ENTRY:
-            raise ValueError(
-                f"ReferenceEntry response_category must be "
-                f"{ResponseCategory.REFERENCE_ENTRY}, got `{self.response_category}`."
-            )
-        return self
-
-    def to_sse_event(self) -> ServerSentEvent:
-        """Convert to SSE event format."""
-        return ServerSentEvent(
-            event="copilot_response",
-            data={
-                "index": self.index,
-                "title": self.title,
-                "url": self.url,
-                "response_category": self.response_category.value,
-                "completeness": self.response_completeness.value,
-            },
-        )
-
-
-class ReferenceSection(CopilotOutput):
-    """Represents a reference section with documentation links."""
-
-    references: list[ReferenceEntry]
-    response_category: ResponseCategory = Field(
-        default=ResponseCategory.REFERENCE,
-        frozen=True,
-    )
-    response_completeness: ResponseCompleteness = ResponseCompleteness.COMPLETE
-
-    @model_validator(mode="after")
-    def validate_response_category(self) -> "ReferenceSection":
-        """Validate that response_category has the correct default value."""
-        if self.response_category != ResponseCategory.REFERENCE:
-            raise ValueError(
-                f"ReferenceSection response_category must be "
-                f"{ResponseCategory.REFERENCE}, got `{self.response_category}`."
-            )
-        return self
-
-    def to_sse_event(self) -> ServerSentEvent:
-        """Convert to SSE event format."""
-        return ServerSentEvent(
-            event="copilot_response",
-            data={
-                "references": [
-                    reference.model_dump(include={"index", "title", "url"})
-                    for reference in self.references
-                ],
-                "response_category": self.response_category.value,
-                "completeness": self.response_completeness.value,
-            },
-        )
-
-    def sort_references(self) -> None:
-        """Sort references by index value."""
-        sorted_references = sorted(
-            self.references, key=lambda reference: (0, int(reference.index))
-        )
-
-        self.references = sorted_references
-
-
-class CopilotContext(BaseModel):
-    """Model containing the context used by the copilot to generate a response."""
-
-    assistant_logs: str = Field("")
-    assistant_files: Dict[str, str] = Field({})
-    copilot_chat_history: List["CopilotChatMessage"] = Field([])
-    tracker_context: Optional[TrackerContext] = Field(None)
-
-    class Config:
-        """Config for LLMBuilderContext."""
-
-        arbitrary_types_allowed = True
-
-
-class UsageStatistics(BaseModel):
-    prompt_tokens: Optional[int] = None
-    completion_tokens: Optional[int] = None
-    total_tokens: Optional[int] = None
-    model: Optional[str] = None
-
-    def reset(self) -> None:
-        """Reset usage statistics to their default values."""
-        self.prompt_tokens = None
-        self.completion_tokens = None
-        self.total_tokens = None
-        self.model = None
-
-    def update_from_stream_chunk(self, chunk: ChatCompletionChunk) -> None:
-        """Update usage statistics from an OpenAI stream chunk.
-
-        Args:
-            chunk: The OpenAI stream chunk containing usage statistics.
-        """
-        if not (usage := getattr(chunk, "usage", None)):
-            return
-
-        self.prompt_tokens = usage.prompt_tokens
-        self.completion_tokens = usage.completion_tokens
-        self.total_tokens = usage.total_tokens
-        self.model = getattr(chunk, "model", None)