rasa-pro 3.14.0.dev6__py3-none-any.whl → 3.14.0.dev8__py3-none-any.whl

rasa/builder/copilot/copilot_response_handler.py

@@ -0,0 +1,522 @@
+import copy
+import re
+from collections import deque
+from typing import AsyncGenerator, Deque, Dict, List, Optional
+
+import structlog
+
+from rasa.builder.copilot.copilot_templated_message_provider import (
+    load_copilot_handler_default_responses,
+)
+from rasa.builder.copilot.exceptions import (
+    CopilotFinalBufferReached,
+    CopilotStreamEndedEarly,
+)
+from rasa.builder.copilot.models import (
+    CopilotOutput,
+    GeneratedContent,
+    ReferenceEntry,
+    ReferenceSection,
+    ResponseCategory,
+    ResponseCompleteness,
+)
+from rasa.builder.document_retrieval.models import Document
+from rasa.builder.guardrails.constants import BLOCK_SCOPE_USER, BlockScope
+
+structlogger = structlog.get_logger()
+
+# Controlled prediction markers
+ROLEPLAY_PREDICTION = "[ROLEPLAY_REQUEST_DETECTED]"
+OUT_OF_SCOPE_PREDICTION = "[OUT_OF_SCOPE_REQUEST_DETECTED]"
+ERROR_FALLBACK_PREDICTION = "[ERROR_FALLBACK]"
+KNOWLEDGE_BASE_ACCESS_REQUESTED_PREDICTION = "[NO_KNOWLEDGE_BASE_ACCESS]"
+
+# Load predefined for controlled predictions from YAML
+_handler_responses = load_copilot_handler_default_responses()
+
+# Prediction marker to response mapping
+PREDICTION_RESPONSES = {
+    ROLEPLAY_PREDICTION: (
+        _handler_responses.get("roleplay_response", ""),
+        ResponseCategory.ROLEPLAY_DETECTION,
+    ),
+    OUT_OF_SCOPE_PREDICTION: (
+        _handler_responses.get("out_of_scope_response", ""),
+        ResponseCategory.OUT_OF_SCOPE_DETECTION,
+    ),
+    ERROR_FALLBACK_PREDICTION: (
+        _handler_responses.get("error_fallback_response", ""),
+        ResponseCategory.ERROR_FALLBACK,
+    ),
+    KNOWLEDGE_BASE_ACCESS_REQUESTED_PREDICTION: (
+        _handler_responses.get("knowledge_base_access_requested_response", ""),
+        ResponseCategory.KNOWLEDGE_BASE_ACCESS_REQUESTED,
+    ),
+}
+
+# Other response constants
+GUARDRAIL_POLICY_VIOLATION_RESPONSE = _handler_responses.get(
+    "guardrail_policy_violation_response", ""
+)
+COPILOT_REDACTED_MESSAGE = _handler_responses.get("copilot_redacted_message", "")
+
+# Guardrails blocked responses
+GUARDRAIL_BLOCKED_USER_RESPONSE = _handler_responses.get(
+    "guardrail_blocked_user_response", ""
+)
+GUARDRAIL_BLOCKED_PROJECT_RESPONSE = _handler_responses.get(
+    "guardrail_blocked_project_response", ""
+)
+
+# Common LLM response prefixes and suffixes before the actual content. These are removed
+# from the content.
+LLM_PREFIXES_TO_SUFFIX_REMOVE = {
+    "```": "```",
+    "```markdown": "```",
+    '"""': '"""',
+}
+
+# Regex pattern to match inline citations in the generated markdown.
+INLINE_CITATION_PATTERN = r"\[([^\]]+)\]\(([^)]+)\)"
+
+
+class CopilotResponseHandler:
+    """Handles the controlled responses from Copilot.
+
+    This handler manages two types of data:
+    - llm_stream_buffer: A list of tokens streamed from the LLM during processing.
+    - generated_responses: A list of cleaned responses.
+
+    Parameters:
+        rolling_buffer_size: Size of the rolling buffer for prefix/suffix handling.
+    """
+
+    def __init__(
+        self,
+        rolling_buffer_size: int = 20,
+    ):
+        self._rolling_buffer_size = rolling_buffer_size
+
+        # Rolling buffer for handling special tokens and prefix/suffix removal.
+        self._rolling_buffer: Deque[str] = deque(maxlen=self._rolling_buffer_size)
+
+        # A list of tokens streamed from the LLM during processing. Tokens are added
+        # to the buffer when the rolling buffer is full.
+        self._llm_stream_buffer: List[str] = []
+
+        # A list of cleaned and generated responses.
+        self._generated_responses: List[GeneratedContent] = []
+
+        # Flag for checking if the stream has ended (raised by the async generator).
+        self._stream_ended: bool = False
+        # Flag for checking if the stream ended before the max expected special response
+        # tokens were reached.
+        self._stream_ended_early: bool = False
+
+        # Maximum number of tokens to check for special responses (e.g. roleplay,
+        # out-of-scope, etc.).
+        self._max_expected_special_response_tokens: int = 20
+
+        # Prefix/suffix tracking
+        self._stream_started_with_prefix: bool = False
+        self._prefix_found: Optional[str] = None
+        self._suffix_found: Optional[str] = None
+
+    @property
+    def llm_stream_buffer(self) -> List[str]:
+        return copy.deepcopy(self._llm_stream_buffer)
+
+    @property
+    def generated_responses(self) -> List[GeneratedContent]:
+        return copy.deepcopy(self._generated_responses)
+
+    @property
+    def llm_stream_buffer_content(self) -> str:
+        return "".join(self._llm_stream_buffer)
+
+    @property
+    def token_count(self) -> int:
+        return len(self._llm_stream_buffer)
+
+    def clear_buffers(self) -> None:
+        """Clear all buffers and reset the handler."""
+        self._rolling_buffer.clear()
+        self._llm_stream_buffer.clear()
+        self._generated_responses.clear()
+        self._stream_started_with_prefix = False
+        self._prefix_found = None
+        self._suffix_found = None
+        self._stream_ended = False
+
+    async def _buffer_stream(
+        self, response_stream: AsyncGenerator[str, None]
+    ) -> AsyncGenerator[str, None]:
+        """Wrapper that performs rolling buffer handling and automatically saves chunks.
+
+        Args:
+            response_stream: The original streaming response from the LLM.
+
+        Yields:
+            The same chunks from the original stream, but persisted in the buffer.
+
+        Raises:
+            StopAsyncIteration: If the stream has ended.
+        """
+        try:
+            while True:
+                # Get the next chunk from LLM stream and add it to the rolling buffer
+                # and the LLM stream buffer.
+                chunk = await response_stream.__anext__()
+                self._rolling_buffer.append(chunk)
+
+                # Only yield when buffer is full to maintain the rolling buffer
+                # behavior
+                if len(self._rolling_buffer) == self._rolling_buffer_size:
+                    # Yield the oldest element.
+                    oldest_element = self._rolling_buffer.popleft()
+                    self._llm_stream_buffer.append(oldest_element)
+                    yield oldest_element
+
+        except StopAsyncIteration:
+            self._llm_stream_buffer.extend(self._rolling_buffer)
+            self._stream_ended = True
+            # Need to raise a new instance of StopAsyncIteration to avoid the
+            # RuntimeError: generator raised StopIteration
+            raise CopilotFinalBufferReached()
+
+        except Exception as e:
+            structlogger.exception(
+                "copilot_response_handler._buffer_stream.unexpected_error",
+            )
+            raise e
+
+    async def handle_response(
+        self, response_stream: AsyncGenerator[str, None]
+    ) -> AsyncGenerator[CopilotOutput, None]:
+        """Intercept a streaming response and handle special responses from the Copilot.
+
+        Args:
+            response_stream: The original streaming response from the LLM.
+
+        Yields:
+            ResponseEvent objects representing either generated tokens, default
+            responses, or reference sections.
+        """
+        # Clear the stream buffer and reference buffer at the start
+        self.clear_buffers()
+
+        try:
+            # Exhaust the buffer early to check for controlled predictions and prefix
+            # detection.
+            await self._exhaust_buffer_for_early_detection(response_stream)
+
+            # Check for controlled predictions in the collected tokens
+            controlled_response = self._check_for_controlled_predictions(
+                self.llm_stream_buffer_content,
+            )
+            if controlled_response is not None:
+                self._generated_responses.append(controlled_response)
+                yield controlled_response
+                return
+
+            # At this point, no controlled predictions were found. Check if the stream
+            # started with a prefix and if present, remove it. Yield the clean content.
+            initial_content = self._remove_prefix(self.llm_stream_buffer_content)
+            if initial_content:
+                generated_content = GeneratedContent(
+                    content=initial_content,
+                    response_category=ResponseCategory.COPILOT,
+                    response_completeness=ResponseCompleteness.TOKEN,
+                )
+                self._generated_responses.append(generated_content)
+                yield generated_content
+
+            # Continue streaming remaining chunks with rolling buffer handling
+            async for chunk in self._buffer_stream(response_stream):
+                generated_content = GeneratedContent(
+                    content=chunk,
+                    response_category=ResponseCategory.COPILOT,
+                    response_completeness=ResponseCompleteness.TOKEN,
+                )
+                self._generated_responses.append(generated_content)
+                yield generated_content
+
+        # Stream ended early
+        except CopilotStreamEndedEarly:
+            # Check for controlled predictions in the collected tokens
+            controlled_response = self._check_for_controlled_predictions(
+                self.llm_stream_buffer_content,
+            )
+            if controlled_response is not None:
+                self._generated_responses.append(controlled_response)
+                yield controlled_response
+                return
+
+            # At this point, no controlled predictions were found. Clean the content
+            # from the prefix and suffix if present.
+            final_content = self._remove_prefix_and_suffix(
+                self.llm_stream_buffer_content,
+            )
+            if final_content:
+                generated_content = GeneratedContent(
+                    content=final_content,
+                    response_category=ResponseCategory.COPILOT,
+                    response_completeness=ResponseCompleteness.COMPLETE,
+                )
+                self._generated_responses.append(generated_content)
+                yield generated_content
+
+        # Stream has ended, process the final rolling buffer. Remove the suffix if
+        # present.
+        except CopilotFinalBufferReached:
+            final_content = self._remove_suffix(self.llm_stream_buffer_content)
+            if final_content:
+                generated_content = GeneratedContent(
+                    content=final_content,
+                    response_category=ResponseCategory.COPILOT,
+                    response_completeness=ResponseCompleteness.COMPLETE,
+                )
+                self._generated_responses.append(generated_content)
+                yield generated_content
+
+        # Unexpected error occurred.
+        except Exception as e:
+            structlogger.exception(
+                "copilot_response_handler.handle_response.unexpected_error",
+            )
+            raise e
+
+    async def _exhaust_buffer_for_early_detection(
+        self, response_stream: AsyncGenerator[str, None]
+    ) -> None:
+        """Exhaust the buffer for early detection.
+
+        Args:
+            response_stream: The original streaming response from the LLM.
+        """
+        try:
+            async for _ in self._buffer_stream(response_stream):
+                if self.token_count >= self._max_expected_special_response_tokens:
+                    break
+        except CopilotFinalBufferReached:
+            # The stream ended before the max expected special response tokens
+            # were reached.
+            raise CopilotStreamEndedEarly()
+
+    def _check_for_controlled_predictions(
+        self, content: str
+    ) -> Optional[CopilotOutput]:
+        """Check for controlled predictions in the collected tokens.
+
+        Returns:
+            Controlled response if found, None otherwise.
+        """
+        # Check for controlled predictions and return appropriate responses
+        for prediction_marker, (response, category) in PREDICTION_RESPONSES.items():
+            if prediction_marker in content:
+                log_message = f"copilot_response_handler.{category.value}_detected"
+                structlogger.info(log_message)
+                return GeneratedContent(
+                    response_category=category,
+                    content=response,
+                    response_completeness=ResponseCompleteness.COMPLETE,
+                )
+
+        return None
+
+    def _remove_prefix(self, content: str) -> str:
+        """Process the initial content from the buffer, handling prefix removal.
+
+        Returns:
+            Processed content with prefix removed if applicable.
+        """
+        # Check if content starts with any of the known prefixes
+        for prefix in LLM_PREFIXES_TO_SUFFIX_REMOVE.keys():
+            if content.startswith(prefix):
+                self._stream_started_with_prefix = True
+                self._prefix_found = prefix
+                self._code_block_depth = 1
+                structlogger.debug(
+                    "copilot_response_handler.handle_response.prefix_detected",
+                    prefix=prefix,
+                    content_length=len(content),
+                )
+                return content[len(prefix) :]
+
+        return content
+
+    def _remove_suffix(self, content: str) -> str:
+        """Process the rolling buffer content, handling suffix removal.
+
+        Returns:
+            Processed content with suffix removed if applicable.
+        """
+        # Check if content ends with any of the known suffixes
+        for suffix in LLM_PREFIXES_TO_SUFFIX_REMOVE.values():
+            if content.endswith(suffix):
+                self._stream_ended_with_suffix = True
+                self._suffix_found = suffix
+                structlogger.debug(
+                    "copilot_response_handler.handle_response.suffix_detected",
+                    suffix=suffix,
+                    content_length=len(content),
+                )
+                return content[: -len(suffix)]
+
+        return content
+
+    def _remove_prefix_and_suffix(self, content: str) -> str:
+        """Remove the prefix and suffix from the content.
+
+        Returns:
+            Processed content with prefix and suffix removed if applicable.
+        """
+        content = self._remove_prefix(content)
+        content = self._remove_suffix(content)
+        return content
+
+    @staticmethod
+    def respond_to_guardrail_policy_violations() -> GeneratedContent:
+        """Respond to guardrail policy violations.
+
+        Returns:
+            CopilotOutput object with the response.
+        """
+        return GeneratedContent(
+            response_category=ResponseCategory.GUARDRAILS_POLICY_VIOLATION,
+            content=GUARDRAIL_POLICY_VIOLATION_RESPONSE,
+            response_completeness=ResponseCompleteness.COMPLETE,
+        )
+
+    @staticmethod
+    def respond_to_guardrail_blocked(scope: BlockScope) -> GeneratedContent:
+        """Return a blocked response for user or project scope.
+
+        Args:
+            scope: 'user' for user-level block, 'project' for project-level block.
+
+        Returns:
+            GeneratedContent with GUARDRAILS_BLOCKED category.
+        """
+        content = (
+            GUARDRAIL_BLOCKED_USER_RESPONSE
+            if scope == BLOCK_SCOPE_USER
+            else GUARDRAIL_BLOCKED_PROJECT_RESPONSE
+        )
+        return GeneratedContent(
+            response_category=ResponseCategory.GUARDRAILS_BLOCKED,
+            content=content,
+            response_completeness=ResponseCompleteness.COMPLETE,
+        )
+
+    @staticmethod
+    def get_copilot_redacted_message() -> str:
+        """Get the redacted message for copilot responses.
+
+        Returns:
+            Redacted message for copilot responses.
+        """
+        return COPILOT_REDACTED_MESSAGE
+
+    def extract_references(self, documents: List[Document]) -> ReferenceSection:
+        """Extract references from the LLM stream buffer content.
+
+        This method performs regex matching to find markdown links in the format:
+        [text](url).
+
+        The matched links are validated against the provided documents, and a
+        ReferenceSection is returned with valid references.
+
+        Note:
+            This method can only be called after `handle_response` has been called,
+            as it relies on the LLM stream buffer content that gets populated
+            during the response handling process.
+
+        Args:
+            documents: List of Document objects to match URLs against
+
+        Returns:
+            ReferenceSection containing reference entries ordered by reference text.
+
+        Raises:
+            RuntimeError: If called before `handle_response` has been called.
+        """
+        if not self.llm_stream_buffer_content:
+            message = (
+                "`extract_references` can only be called after `handle_response` "
+                "has been called to populate the LLM stream buffer."
+            )
+            structlogger.error(
+                "copilot_response_handler.extract_references.buffer_not_populated",
+                event_info=message,
+                llm_stream_buffer_content=self.llm_stream_buffer_content,
+            )
+            raise RuntimeError(message)
+
+        # Find all matches in the buffered content
+        matches = re.findall(INLINE_CITATION_PATTERN, self.llm_stream_buffer_content)
+
+        # Create document lookup for O(1) access
+        document_urls_to_documents: Dict[str, Document] = {
+            document.url: document for document in documents if document.url
+        }
+
+        # Use regular dict to collect references, keyed by reference_text
+        used_references: Dict[str, ReferenceEntry] = {}
+
+        for reference_text, reference_url in matches:
+            # Validate reference text format
+            if not reference_text.isdigit():
+                structlogger.warning(
+                    "copilot_response_handler"
+                    ".extract_references."
+                    "invalid_reference_number",
+                    event_info="Reference text is not in expected number format.",
+                    reference_text=reference_text,
+                    reference_url=reference_url,
+                )
+                continue
+
+            # Check if URL exists in documents
+            if reference_url not in document_urls_to_documents:
+                structlogger.warning(
+                    "copilot_response_handler.extract_references.url_not_found",
+                    event_info=(
+                        "URL not found in provided documents. Omitted from reference "
+                        "section.",
+                    ),
+                    reference_url=reference_url,
+                    available_urls=list(document_urls_to_documents.keys()),
+                )
+                continue
+
+            # Check for duplicate reference text
+            # (same reference number used multiple times)
+            if reference_text in used_references:
+                existing_entry = used_references[reference_text]
+                if existing_entry.url != reference_url:
+                    structlogger.warning(
+                        "copilot_response_handler.extract_references.duplicate_reference_text",
+                        event_info=(
+                            "Same reference text used for different URLs. "
+                            "Keeping first occurrence."
+                        ),
+                        reference_text=reference_text,
+                        first_url=existing_entry.url,
+                        second_url=reference_url,
+                    )
+                # Skip this duplicate reference text
+                continue
+
+            # Create reference entry
+            document = document_urls_to_documents[reference_url]
+            used_references[reference_text] = ReferenceEntry(
+                index=int(reference_text),
+                title=document.title or f"Reference {reference_text}",
+                url=reference_url,
+            )
+
+        # Create and sort the reference section
+        reference_section = ReferenceSection(references=list(used_references.values()))
+        reference_section.sort_references()
+        return reference_section

rasa/builder/copilot/copilot_templated_message_provider.py

@@ -0,0 +1,58 @@
+import importlib.resources
+from typing import Dict
+
+import structlog
+import yaml  # type: ignore
+
+from rasa.builder.copilot.constants import (
+    COPILOT_HANDLER_RESPONSES_FILE,
+    COPILOT_MESSAGE_TEMPLATES_DIR,
+    RASA_INTERNAL_MESSAGES_TEMPLATES_FILE,
+)
+from rasa.shared.constants import PACKAGE_NAME
+
+structlogger = structlog.get_logger()
+
+
+def load_copilot_internal_message_templates() -> Dict[str, str]:
+    """Load internal message templates from the YAML configuration file.
+
+    Returns:
+        Dictionary mapping template names to template text.
+    """
+    try:
+        config = yaml.safe_load(
+            importlib.resources.read_text(
+                f"{PACKAGE_NAME}.{COPILOT_MESSAGE_TEMPLATES_DIR}",
+                RASA_INTERNAL_MESSAGES_TEMPLATES_FILE,
+            )
+        )
+        return config.get("templates", {})
+    except Exception as e:
+        structlogger.error(
+            "copilot_templated_message_provider.failed_to_load_templates",
+            error=e,
+        )
+        return dict()
+
+
+def load_copilot_handler_default_responses() -> Dict[str, str]:
+    """Load handler responses from the YAML configuration file.
+
+    Returns:
+        Dictionary mapping response names to response text.
+    """
+    try:
+        config = yaml.safe_load(
+            importlib.resources.read_text(
+                f"{PACKAGE_NAME}.{COPILOT_MESSAGE_TEMPLATES_DIR}",
+                COPILOT_HANDLER_RESPONSES_FILE,
+            )
+        )
+        return config.get("responses", {})
+    except Exception as e:
+        structlogger.error(
+            "copilot_response_handler.failed_to_load_responses",
+            error=e,
+        )
+        return dict()

rasa/builder/copilot/exceptions.py

@@ -0,0 +1,32 @@
+class CopilotFinalBufferReached(Exception):
+    """Raised when the rolling buffer ends."""
+
+    pass
+
+
+class CopilotStreamEndedEarly(Exception):
+    """Raised when the stream ended early.
+
+    This happens when the stream ended before the max expected special response
+    tokens were reached.
+    """
+
+    pass
+
+
+class CopilotStreamError(Exception):
+    """Raised when the stream fails."""
+
+    pass
+
+
+class InvalidCopilotChatHistorySignature(Exception):
+    """Raised when the provided history signature does not match."""
+
+    pass
+
+
+class MissingCopilotChatHistorySignature(Exception):
+    """Raised when a required history signature is missing."""
+
+    pass