rasa-pro 3.14.0rc4__py3-none-any.whl → 3.15.0a1__py3-none-any.whl

rasa/builder/service.py

@@ -5,6 +5,7 @@ import time
 from http import HTTPStatus
 from typing import Any, Optional
 
+import langfuse
 import structlog
 from sanic import Blueprint, HTTPResponse, response
 from sanic.request import Request
@@ -41,7 +42,6 @@ from rasa.builder.copilot.signing import (
     create_signature_envelope_for_text,
     verify_signature,
 )
-from rasa.builder.copilot.telemetry import CopilotTelemetry
 from rasa.builder.download import create_bot_project_archive
 from rasa.builder.guardrails.constants import (
     BLOCK_SCOPE_PROJECT,
@@ -65,6 +65,7 @@ from rasa.builder.models import (
     ApiErrorResponse,
     AssistantInfo,
     BotData,
+    BotFiles,
     JobCreateResponse,
     JobStatus,
     JobStatusEvent,
@@ -74,6 +75,8 @@ from rasa.builder.models import (
 )
 from rasa.builder.project_generator import ProjectGenerator
 from rasa.builder.shared.tracker_context import TrackerContext
+from rasa.builder.telemetry.copilot_langfuse_telemetry import CopilotLangfuseTelemetry
+from rasa.builder.telemetry.copilot_segment_telemetry import CopilotSegmentTelemetry
 from rasa.core.agent import Agent
 from rasa.core.channels.studio_chat import StudioChatInput
 from rasa.core.exceptions import AgentNotReady
@@ -1020,6 +1023,9 @@ async def download_bot_project(request: Request) -> HTTPResponse:
     schema=str,
 )
 @protected()
+# Disable automatic input/output capture for langfuse tracing
+# This allows manual control over what data is sent to langfuse
+@langfuse.observe(capture_input=False, capture_output=False)
 async def copilot(request: Request) -> None:
     """Handle copilot requests with streaming markdown responses."""
     sse = await request.respond(content_type="text/event-stream")
@@ -1046,9 +1052,12 @@ async def copilot(request: Request) -> None:
             )
             return
 
-        telemetry = CopilotTelemetry(project_id=HELLO_RASA_PROJECT_ID, user_id=user_id)
+        telemetry = CopilotSegmentTelemetry(
+            project_id=HELLO_RASA_PROJECT_ID, user_id=user_id
+        )
         structlogger.debug("builder.copilot.telemetry.request.init")
 
+        # TODO: This can be removed once Langfuse is completed.
         if req.last_message and req.last_message.role == ROLE_USER:
             structlogger.debug("builder.copilot.telemetry.request.user_turn")
             # Offload telemetry logging to a background task
@@ -1088,26 +1097,9 @@ async def copilot(request: Request) -> None:
             return
 
         # 4. Get the necessary context for the copilot
-        tracker = await current_tracker_from_input_channel(request.app, req.session_id)
-        tracker_context = TrackerContext.from_tracker(
-            tracker, max_turns=COPILOT_ASSISTANT_TRACKER_MAX_TURNS
-        )
-        if (
-            tracker_context is not None
-            and llm_service.guardrails_policy_checker is not None
-        ):
-            tracker_context = await llm_service.guardrails_policy_checker.check_assistant_chat_for_policy_violations(  # noqa: E501
-                tracker_context=tracker_context,
-                hello_rasa_user_id=user_id,
-                hello_rasa_project_id=HELLO_RASA_PROJECT_ID,
-                lakera_project_id=LAKERA_ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
-            )
-
-        # Copilot doesn't need to know about the docs and any file that is not a core
-        # assistant file
-        relevant_assistant_files = project_generator.get_bot_files(
-            exclude_docs_directory=True,
-            allowed_file_extensions=["yaml", "yml", "py", "jinja", "jinja2"],
+        tracker_context = await get_tracker_context_for_copilot(request, req, user_id)
+        relevant_assistant_files = get_relevant_assistant_files_for_copilot(
+            project_generator,
         )
         context = CopilotContext(
             tracker_context=tracker_context,
@@ -1162,7 +1154,7 @@ async def copilot(request: Request) -> None:
         async for token in intercepted_stream:
             await sse.send(token.to_sse_event().format())
 
-        # 8. Offload telemetry logging to a background task
+        # 8a. Offload metabase telemetry logging to a background task
         request.app.add_task(
             asyncio.to_thread(
                 telemetry.log_copilot_from_handler,
@@ -1177,9 +1169,27 @@ async def copilot(request: Request) -> None:
                     else None
                 ),
                 tracker_event_attachments=generation_context.tracker_event_attachments,
-                **copilot_client.usage_statistics.model_dump(),
+                model=copilot_client.usage_statistics.model or "N/A",
+                prompt_tokens=copilot_client.usage_statistics.prompt_tokens or 0,
+                cached_prompt_tokens=(
+                    copilot_client.usage_statistics.cached_prompt_tokens or 0
+                ),
+                completion_tokens=(
+                    copilot_client.usage_statistics.completion_tokens or 0
+                ),
+                total_tokens=copilot_client.usage_statistics.total_tokens or 0,
             )
         )
+        # 8b. Setup output trace attributes for Langfuse
+        CopilotLangfuseTelemetry.setup_copilot_endpoint_call_trace_attributes(
+            hello_rasa_project_id=HELLO_RASA_PROJECT_ID or "N/A",
+            chat_id=req.session_id or "N/A",
+            user_id=user_id,
+            request=req,
+            handler=copilot_response_handler,
+            relevant_documents=generation_context.relevant_documents,
+            copilot_context=context,
+        )
 
         # 9. Once the stream is over, extract and send references
         #    if any documents were used
@@ -1365,3 +1375,70 @@ async def _handle_guardrail_violation_and_maybe_block(
 
     await sse.send(message.to_sse_event().format())
     return message
+
+
+@langfuse.observe(capture_input=False, capture_output=False)
+async def get_tracker_context_for_copilot(
+    request: Request,
+    req: CopilotRequest,
+    user_id: str,
+) -> Optional[TrackerContext]:
+    """Check the assistant chat for guardrail policy violations.
+
+    Args:
+        request: The request object.
+        req: The CopilotRequest object.
+        user_id: The user ID.
+
+    Returns:
+        The tracker context if the tracker is available.
+    """
+    tracker = await current_tracker_from_input_channel(request.app, req.session_id)
+    tracker_context = TrackerContext.from_tracker(
+        tracker, max_turns=COPILOT_ASSISTANT_TRACKER_MAX_TURNS
+    )
+    if (
+        tracker_context is not None
+        and llm_service.guardrails_policy_checker is not None
+    ):
+        tracker_context = await llm_service.guardrails_policy_checker.check_assistant_chat_for_policy_violations(  # noqa: E501
+            tracker_context=tracker_context,
+            hello_rasa_user_id=user_id,
+            hello_rasa_project_id=HELLO_RASA_PROJECT_ID,
+            lakera_project_id=LAKERA_ASSISTANT_HISTORY_GUARDRAIL_PROJECT_ID,
+        )
+
+    # Track the retrieved tracker context
+    CopilotLangfuseTelemetry.trace_copilot_tracker_context(
+        tracker_context=tracker_context,
+        max_conversation_turns=COPILOT_ASSISTANT_TRACKER_MAX_TURNS,
+        session_id=req.session_id,
+    )
+
+    return tracker_context
+
+
+@langfuse.observe(capture_input=False, capture_output=False)
+def get_relevant_assistant_files_for_copilot(
+    project_generator: ProjectGenerator,
+) -> BotFiles:
+    """Get the relevant assistant files for the copilot.
+
+    Args:
+        project_generator: The project generator.
+
+    Returns:
+        The relevant assistant files.
+    """
+    # Copilot doesn't need to know about the docs and any file that is not a core
+    # assistant file
+    files = project_generator.get_bot_files(
+        exclude_docs_directory=True,
+        allowed_file_extensions=["yaml", "yml", "py", "jinja", "jinja2"],
+    )
+
+    # Track the retrieved assistant files
+    CopilotLangfuseTelemetry.trace_copilot_relevant_assistant_files(
+        relevant_assistant_files=files,
+    )
+    return files

rasa/builder/telemetry/copilot_langfuse_telemetry.py

@@ -0,0 +1,384 @@
+from functools import wraps
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Callable,
+    Dict,
+    List,
+    Optional,
+)
+
+if TYPE_CHECKING:
+    from rasa.builder.copilot.copilot import Copilot
+    from rasa.builder.copilot.models import CopilotContext
+    from rasa.builder.document_retrieval.inkeep_document_retrieval import (
+        InKeepDocumentRetrieval,
+    )
+
+import langfuse
+import structlog
+
+from rasa.builder.copilot.copilot_response_handler import CopilotResponseHandler
+from rasa.builder.copilot.models import (
+    CopilotRequest,
+    EventContent,
+    UsageStatistics,
+    UserChatMessage,
+)
+from rasa.builder.document_retrieval.models import Document
+from rasa.builder.models import BotFiles
+from rasa.builder.shared.tracker_context import TrackerContext
+
+structlogger = structlog.get_logger()
+
+
+class CopilotLangfuseTelemetry:
+    @staticmethod
+    def trace_copilot_tracker_context(
+        tracker_context: Optional[TrackerContext],
+        max_conversation_turns: int,
+        session_id: str,
+    ) -> None:
+        """Trace the copilot tracker context.
+
+        Args:
+            tracker_context: The tracker context.
+            max_conversation_turns: The maximum number of conversation turns to be
+                fetched from the tracker.
+            session_id: The session ID used to fetch the right tracker.
+        """
+        langfuse_client = langfuse.get_client()
+        # Use `update_current_span` to update the current span of the trace.
+        langfuse_client.update_current_span(
+            output={
+                "tracker_context": (
+                    tracker_context.model_dump() if tracker_context else None
+                ),
+            },
+            metadata={
+                "max_conversation_turns": max_conversation_turns,
+                "session_id": session_id,
+            },
+        )
+
+    @staticmethod
+    def trace_copilot_relevant_assistant_files(
+        relevant_assistant_files: BotFiles,
+    ) -> None:
+        """Trace the copilot relevant assistant files.
+
+        Args:
+            relevant_assistant_files: The relevant assistant files.
+        """
+        langfuse_client = langfuse.get_client()
+        # Use `update_current_span` to update the current span of the trace.
+        langfuse_client.update_current_span(
+            output={
+                "relevant_assistant_files": relevant_assistant_files,
+            },
+        )
+
+    @staticmethod
+    def setup_copilot_endpoint_call_trace_attributes(
+        hello_rasa_project_id: str,
+        chat_id: str,
+        user_id: str,
+        request: CopilotRequest,
+        handler: CopilotResponseHandler,
+        relevant_documents: list[Document],
+        copilot_context: "CopilotContext",
+    ) -> None:
+        """Set up the current langfuse trace with project and user context.
+
+        Args:
+            hello_rasa_project_id: The Hello Rasa project ID.
+            chat_id: The chat/conversation ID.
+            user_id: The user ID.
+            request: The parsed CopilotRequest object.
+            handler: The response handler containing generated responses.
+            relevant_documents: The relevant documents used to generate the response.
+        """
+        langfuse_client = langfuse.get_client()
+        user_message = CopilotLangfuseTelemetry._extract_last_user_message_content(
+            request
+        )
+        tracker_event_attachments = (
+            CopilotLangfuseTelemetry._extract_tracker_event_attachments(request)
+        )
+        response_category = CopilotLangfuseTelemetry._extract_response_category(handler)
+        reference_section_entries = CopilotLangfuseTelemetry._extract_references(
+            handler, relevant_documents
+        )
+
+        # Create a session ID as a composite ID from project id, user id and chat id
+        session_id = CopilotLangfuseTelemetry._create_session_id(
+            hello_rasa_project_id, user_id, chat_id
+        )
+        # Use `update_current_trace` to update the top level trace.
+        langfuse_client.update_current_trace(
+            user_id=user_id,
+            session_id=session_id,
+            input={
+                "message": user_message,
+                "tracker_event_attachments": tracker_event_attachments,
+            },
+            output={
+                "answer": CopilotLangfuseTelemetry._full_text(handler),
+                "response_category": response_category,
+                "references": reference_section_entries,
+            },
+            metadata={
+                "ids": {
+                    "user_id": user_id,
+                    "project_id": hello_rasa_project_id,
+                    "chat_history_id": chat_id,
+                },
+                "copilot_additional_context": {
+                    "relevant_documents": [
+                        doc.model_dump() for doc in relevant_documents
+                    ],
+                    "relevant_assistant_files": copilot_context.assistant_files,
+                    "assistant_tracker_context": (
+                        copilot_context.tracker_context.model_dump()
+                        if copilot_context.tracker_context
+                        else None
+                    ),
+                    "assistant_logs": copilot_context.assistant_logs,
+                    "copilot_chat_history": [
+                        message.model_dump()
+                        for message in copilot_context.copilot_chat_history
+                    ],
+                },
+            },
+            tags=[response_category] if response_category else [],
+        )
+
+    @staticmethod
+    def trace_copilot_streaming_generation(
+        func: Callable[..., AsyncGenerator[str, None]],
+    ) -> Callable[..., AsyncGenerator[str, None]]:
+        """Custom decorator for tracing async streaming of the Copilot's LLM generation.
+
+        This decorator handles Langfuse tracing for async streaming of the Copilot's LLM
+        generation by manually managing the generation span and updating it with usage
+        statistics after the stream completes.
+        """
+
+        @wraps(func)
+        async def wrapper(
+            self: "Copilot", messages: List[Dict[str, Any]]
+        ) -> AsyncGenerator[str, None]:
+            langfuse_client = langfuse.get_client()
+
+            with langfuse_client.start_as_current_generation(
+                name=f"{self.__class__.__name__}.{func.__name__}",
+                input={"messages": messages},
+            ) as generation:
+                output = []
+                # Call the original streaming function and start capturing the output
+                async for chunk in func(self, messages):
+                    output.append(chunk)
+                    yield chunk
+
+                # Update the span's model parameters and output after streaming is
+                # complete
+                generation.update(
+                    model_parameters=self.llm_config, output="".join(output)
+                )
+
+                # Update the span's usage statistics after streaming is complete
+                if self.usage_statistics:
+                    CopilotLangfuseTelemetry._update_generation_span_with_usage_statistics(
+                        generation, self.usage_statistics
+                    )
+
+        return wrapper
+
+    @staticmethod
+    def trace_document_retrieval_generation(
+        func: Callable[..., Any],
+    ) -> Callable[..., Any]:
+        """Custom decorator for tracing document retrieval generation with Langfuse.
+
+        This decorator handles Langfuse tracing for document retrieval API calls
+        by manually managing the generation span and updating it with usage statistics.
+        """
+
+        @wraps(func)
+        async def wrapper(
+            self: "InKeepDocumentRetrieval",
+            query: str,
+            temperature: float,
+            timeout: float,
+        ) -> Any:
+            langfuse_client = langfuse.get_client()
+
+            with langfuse_client.start_as_current_generation(
+                name=f"{self.__class__.__name__}.{func.__name__}",
+                input={
+                    "query": query,
+                    "temperature": temperature,
+                    "timeout": timeout,
+                },
+            ) as generation:
+                # Call the original function
+                response = await func(self, query, temperature, timeout)
+
+                # Update the span with response content
+                generation.update(
+                    output=response,
+                    model_parameters={
+                        "temperature": str(temperature),
+                        "timeout": str(timeout),
+                    },
+                )
+
+                # Update usage statistics if available
+                usage_statistics = UsageStatistics.from_chat_completion_response(
+                    response
+                )
+                if usage_statistics:
+                    CopilotLangfuseTelemetry._update_generation_span_with_usage_statistics(
+                        generation, usage_statistics
+                    )
+
+                return response
+
+        return wrapper
+
+    @staticmethod
+    def _extract_last_user_message_content(request: CopilotRequest) -> Optional[str]:
+        """Extract the last user message from the CopilotRequest object.
+
+        Args:
+            request: The CopilotRequest object.
+        """
+        if not isinstance(request.last_message, UserChatMessage):
+            return None
+        return request.last_message.get_flattened_text_content()
+
+    @staticmethod
+    def _extract_tracker_event_attachments(
+        request: CopilotRequest,
+    ) -> list[Dict[str, Any]]:
+        """Extract tracker event attachments from the last user message.
+
+        Args:
+            request: The CopilotRequest object.
+
+        Returns:
+            The event content block sent with the last user message in the
+            dictionary format.
+        """
+        last_message = request.last_message
+        if not isinstance(last_message, UserChatMessage):
+            return []
+        return [
+            attachment.model_dump()
+            for attachment in last_message.get_content_blocks_by_type(EventContent)
+        ]
+
+    @staticmethod
+    def _extract_response_category(handler: CopilotResponseHandler) -> Optional[str]:
+        """Extract the response category from the response handler.
+
+        Args:
+            handler: The response handler containing generated response.
+
+        Returns:
+            The response category of the first generated response, or None if no
+            responses.
+        """
+        if not handler.generated_responses:
+            return None
+        # The handler contains multiple chunks of one response. We use the first chunk's
+        # response category.
+        return handler.generated_responses[0].response_category.value
+
+    @staticmethod
+    def _full_text(handler: CopilotResponseHandler) -> str:
+        """Extract full text from the response handler.
+
+        Args:
+            handler: The response handler containing generated responses.
+
+        Returns:
+            The concatenated content of all generated responses.
+        """
+        return "".join(
+            response.content
+            for response in handler.generated_responses
+            if getattr(response, "content", None)
+        )
+
+    @staticmethod
+    def _extract_references(
+        handler: CopilotResponseHandler,
+        relevant_documents: list[Document],
+    ) -> List[Dict[str, Any]]:
+        """Extract reference entries from the response handler.
+
+        Args:
+            handler: The response handler containing generated responses.
+            relevant_documents: The relevant documents used to generate the response.
+
+        Returns:
+            A list of reference entries in dictionary format.
+        """
+        if not relevant_documents:
+            return []
+
+        reference_entries: list[Dict[str, Any]] = []
+        reference_section = handler.extract_references(relevant_documents)
+        for reference_entry in reference_section.references:
+            reference_entries.append(
+                reference_entry.model_dump(
+                    exclude={"response_category", "response_completeness"}
+                )
+            )
+
+        return reference_entries
+
+    @staticmethod
+    def _update_generation_span_with_usage_statistics(
+        generation_span: langfuse.LangfuseGeneration,
+        usage_statistics: UsageStatistics,
+    ) -> None:
+        """Update the generation span with the usage statistics.
+
+        Args:
+            generation_span: The generation span.
+            usage_statistics: The usage statistics of the generation.
+        """
+        generation_span.update(
+            usage_details={
+                "input_non_cached_usage": (
+                    usage_statistics.non_cached_prompt_tokens or 0
+                ),
+                "input_cached_usage": usage_statistics.cached_prompt_tokens or 0,
+                "output_usage": usage_statistics.completion_tokens or 0,
+                "total": usage_statistics.total_tokens or 0,
+            },
+            cost_details={
+                "input_non_cached_cost": usage_statistics.non_cached_cost or 0,
+                "input_cached_cost": usage_statistics.cached_cost or 0,
+                "output_cost": usage_statistics.output_cost or 0,
+                "total": usage_statistics.total_cost or 0,
+            },
+            model=usage_statistics.model,
+        )
+
+    @staticmethod
+    def _create_session_id(
+        hello_rasa_project_id: str,
+        user_id: str,
+        chat_id: str,
+    ) -> str:
+        """Create a session ID as a composite from project id, user id and chat id."""
+        pattern = "PID-{project_id}-UID-{user_id}-CID-{chat_id}"
+        return pattern.format(
+            project_id=hello_rasa_project_id,
+            user_id=user_id,
+            chat_id=chat_id,
+        )

rasa/builder/{copilot/telemetry.py → telemetry/copilot_segment_telemetry.py}

@@ -1,14 +1,25 @@
 import datetime as dt
 import os
 import uuid
-from typing import Any, Iterable, Optional, Sequence
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Iterable,
+    Optional,
+    Sequence,
+)
+
+if TYPE_CHECKING:
+    pass
 
 import structlog
 
 from rasa import telemetry
 from rasa.builder.copilot.constants import COPILOT_SEGMENT_WRITE_KEY_ENV_VAR
 from rasa.builder.copilot.copilot_response_handler import CopilotResponseHandler
-from rasa.builder.copilot.models import EventContent
+from rasa.builder.copilot.models import (
+    EventContent,
+)
 from rasa.builder.document_retrieval.models import Document
 from rasa.telemetry import (
     SEGMENT_TRACK_ENDPOINT,
@@ -56,7 +67,7 @@ def _track(event: str, user_id: str, properties: dict) -> None:
         structlogger.warning("builder.telemetry.track_failed", error=str(e))
 
 
-class CopilotTelemetry:
+class CopilotSegmentTelemetry:
     def __init__(
         self,
         *,
@@ -96,6 +107,7 @@ class CopilotTelemetry:
         latency_ms: int,
         model: str,
         input_tokens: Optional[int] = None,
+        cached_prompt_tokens: Optional[int] = None,
         output_tokens: Optional[int] = None,
         total_tokens: Optional[int] = None,
         system_message: Optional[dict[str, Any]] = None,
@@ -112,6 +124,7 @@ class CopilotTelemetry:
             latency_ms: End-to-end Copilot latency to produce this response.
             model: The model used to generate the response.
             input_tokens: Number of input tokens used (optional).
+            cached_prompt_tokens: Number of cached prompt tokens.
             output_tokens: Number of output tokens generated (optional).
             total_tokens: Total number of tokens used (input + output) (optional).
             system_message: The system message used (optional).
@@ -135,6 +148,7 @@ class CopilotTelemetry:
             "latency_ms": latency_ms,
             "model": model,
             "input_tokens": input_tokens,
+            "cached_prompt_tokens": cached_prompt_tokens,
             "output_tokens": output_tokens,
             "total_tokens": total_tokens,
             "chat_history": chat_history,
@@ -193,6 +207,7 @@ class CopilotTelemetry:
         latency_ms: int,
         model: str,
         prompt_tokens: int,
+        cached_prompt_tokens: int,
         completion_tokens: int,
         total_tokens: int,
         system_message: dict[str, Any],
@@ -208,11 +223,13 @@ class CopilotTelemetry:
             latency_ms: End-to-end Copilot latency to produce this response.
             model: The model used to generate the response.
             prompt_tokens: Number of input tokens used.
+            cached_prompt_tokens: Number of cached prompt tokens.
             completion_tokens: Number of output tokens generated.
             total_tokens: Total number of tokens used (input + output).
             system_message: The system message used.
             chat_history: The chat history messages used.
             last_user_message: The last user message used.
+            tracker_event_attachments: List of tracker event attachments.
         """
         structlogger.debug("builder.telemetry.log_copilot_from_handler")
         text = self._full_text(handler)
@@ -223,6 +240,7 @@ class CopilotTelemetry:
             latency_ms=latency_ms,
             model=model,
             input_tokens=prompt_tokens,
+            cached_prompt_tokens=cached_prompt_tokens,
             output_tokens=completion_tokens,
             total_tokens=total_tokens,
             system_message=system_message,

rasa/builder/validation_service.py

@@ -47,6 +47,10 @@ async def validate_project(importer: TrainingDataImporter) -> Optional[str]:
     with capture_validation_logs() as captured_logs:
         try:
             with _mock_sys_exit() as exit_tracker:
+                from rasa.core.config.configuration import Configuration
+
+                Configuration.initialise_empty()
+
                 validate_files(
                     fail_on_warnings=config.VALIDATION_FAIL_ON_WARNINGS,
                     max_history=config.VALIDATION_MAX_HISTORY,