rasa-pro 3.13.1a18__py3-none-any.whl → 3.13.1a20__py3-none-any.whl

rasa/builder/copilot/models.py

@@ -0,0 +1,344 @@
+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 rasa.builder.copilot.constants import ROLE_ASSISTANT, ROLE_COPILOT, 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 = "copilot"
+    ROLEPLAY_DETECTION = "roleplay_detection"
+    OUT_OF_SCOPE_DETECTION = "out_of_scope_detection"
+    GUARDRAILS_POLICY_VIOLATION = "guardrails_policy_violation"
+    REFERENCE = "reference"
+    REFERENCE_ENTRY = "reference_entry"
+
+
+class BaseContent(BaseModel):
+    type: 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 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"
+
+
+class CopilotChatMessage(BaseModel):
+    """Model for a single chat messages between the user and the copilot."""
+
+    role: str = Field(..., pattern=f"^({ROLE_USER}|{ROLE_COPILOT})$")
+    content: List[Union[TextContent, CodeContent, FileContent, EventContent]] = Field(
+        ...,
+        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."
+        ),
+    )
+    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.
+        """
+        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 to_openai_format(self) -> Dict[str, Any]:
+        """Convert to OpenAI message format for API calls."""
+        role = self._map_role_to_openai()
+
+        if role == ROLE_USER:
+            content = self.get_text_content()
+        else:
+            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,
+        }
+        if self.role not in [ROLE_USER, ROLE_COPILOT]:
+            structlogger.error(
+                "copilot_chat_message.to_openai_format.invalid_role",
+                event_info=(
+                    f"Invalid role: `{self.role}`. "
+                    "Only `user` and `copilot` 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,
+            },
+        )
+
+
+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)