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

rasa/builder/copilot/copilot_response_handler.py

@@ -0,0 +1,471 @@
+import copy
+import re
+from collections import deque
+from typing import AsyncGenerator, Deque, List, Optional, Set
+
+import structlog
+
+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
+
+structlogger = structlog.get_logger()
+
+# Controlled prediction markers for detecting roleplay and out-of-scope requests
+ROLEPLAY_PREDICTION = "[ROLEPLAY_REQUEST_DETECTED]"
+OUT_OF_SCOPE_PREDICTION = "[OUT_OF_SCOPE_REQUEST_DETECTED]"
+
+# Predefined responses for controlled predictions
+ROLEPLAY_RESPONSE = """I'm here to help you build and develop your Rasa chatbot, not to
+roleplay as the assistant you're creating.
+
+If you'd like to test how your bot would respond, you can:
+- Share your bot's configuration files and I can help you understand its behavior
+- Ask me to help you modify responses or flows
+- Get guidance on how to improve your bot's conversation flow
+
+What would you like to work on with your Rasa chatbot today?"""
+
+OUT_OF_SCOPE_RESPONSE = """I'm specifically designed to help with Rasa chatbot
+development and customization.
+
+If you have questions about Rasa development or need help with your chatbot, I'm here to
+help! What would you like to work on?"""
+
+# Response for guardrail policy violations
+GUARDRAIL_POLICY_VIOLATION_RESPONSE = """I'm here to help you with Rasa chatbot
+development in a safe and respectful environment. I cannot assist with requests that
+involve harmful content, inappropriate material, or attempts to manipulate my responses.
+
+Please note that continued attempts to use inappropriate content or manipulate
+responses will result in project termination and account blocking.
+"""
+
+# Redacted message for copilot responses to guardrail policy violations.
+COPILOT_REDACTED_MESSAGE = "[Copilot response redacted for guardrails check.]"
+
+# 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_prexif(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.
+        """
+        # If the Copilot detected a roleplay request, return the default response.
+        if ROLEPLAY_PREDICTION in content:
+            structlogger.info("response_interceptor.roleplay_request_detected")
+            return GeneratedContent(
+                response_category=ResponseCategory.ROLEPLAY_DETECTION,
+                content=ROLEPLAY_RESPONSE,
+                response_completeness=ResponseCompleteness.COMPLETE,
+            )
+
+        # If the Copilot detected an out-of-scope request, return the default response.
+        elif OUT_OF_SCOPE_PREDICTION in content:
+            structlogger.info("response_interceptor.out_of_scope_request_detected")
+            return GeneratedContent(
+                response_category=ResponseCategory.OUT_OF_SCOPE_DETECTION,
+                content=OUT_OF_SCOPE_RESPONSE,
+                response_completeness=ResponseCompleteness.COMPLETE,
+            )
+
+        return None
+
+    def _remove_prexif(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_prexif(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 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 valid reference entries.
+
+        Raises:
+            RuntimeError: If called before `handle_response` has populated the buffer.
+        """
+        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)
+
+        valid_references: List[ReferenceEntry] = []
+        document_urls: Set[str] = {
+            document.url for document in documents if document.url
+        }
+
+        for reference_text, reference_url in matches:
+            # Check if link_text is a number
+            if not reference_text.isdigit():
+                structlogger.warning(
+                    "copilot_response_handler.extract_references"
+                    ".link_text_not_a_number",
+                    event_info="Link text is not a number.",
+                    link_text=reference_text,
+                    link_url=reference_url,
+                )
+
+            # Check if URL exists in the documents
+            if reference_url not in document_urls:
+                structlogger.warning(
+                    "copilot_response_handler.extract_references.url_not_found",
+                    event_info=(
+                        "URL not found in provided documents. Omitted from reference "
+                        "section.",
+                    ),
+                    link_url=reference_url,
+                    document_urls=document_urls,
+                )
+                continue
+
+            # Find the corresponding document for this URL and get its index in the list
+            for document_index, document in enumerate(documents, start=1):
+                if document.url == reference_url:
+                    reference_entry = ReferenceEntry(
+                        index=document_index,
+                        title=document.title or f"Reference {document_index}",
+                        url=reference_url,
+                    )
+                    valid_references.append(reference_entry)
+                    break
+
+        return ReferenceSection(references=valid_references)

rasa/builder/copilot/exceptions.py

@@ -0,0 +1,20 @@
+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