Refactor with strands

backend/app/strands_integration/chat_strands.py

@@ -30,6 +30,7 @@
 )
 from .handlers import ToolResultCapture, create_callback_handler
 from .processors import post_process_strands_result
+from .telemetry import StrandsTelemetryManager
 
 logger = logging.getLogger(__name__)
 
@@ -43,6 +44,36 @@ def chat_with_strands(
     on_tool_result: Callable[[ToolRunResult], None] | None = None,
     on_reasoning: Callable[[str], None] | None = None,
 ) -> tuple[ConversationModel, MessageModel]:
+    """
+    Chat with Strands agents.
+
+    Architecture Overview:
+
+    1. Reasoning Content:
+       - Streaming: CallbackHandler processes reasoning events for real-time display
+       - Persistence: Telemetry (ReasoningSpanProcessor) extracts from OpenTelemetry spans
+
+    2. Tool Use/Result (Thinking Log):
+       - Streaming: ToolResultCapture processes tool events for real-time display
+       - Persistence: ToolResultCapture stores processed data for DynamoDB storage
+
+    3. Related Documents (Citations):
+       - Source: ToolResultCapture only
+       - Reason: Requires access to raw tool results for source_link extraction
+
+    Why This Hybrid Approach:
+
+    - ToolResultCapture: Processes raw tool results during execution hooks, enabling
+      source_link extraction and citation functionality. Telemetry only captures
+      post-processed data, losing metadata required for citations.
+
+    - Telemetry: Captures complete reasoning content from OpenTelemetry spans,
+      providing reliable persistence for reasoning data that may not be available
+      in final AgentResult when tools are used.
+
+    - CallbackHandler: Handles real-time streaming of reasoning content during
+      agent execution for immediate user feedback.
+    """
     user_msg_id, conversation, bot = prepare_conversation(user, chat_input)
 
     display_citation = bot is not None and bot.display_retrieved_chunks
@@ -88,6 +119,10 @@ def chat_with_strands(
 
     continue_generate = chat_input.continue_generate
 
+    # Setup telemetry manager for reasoning capture
+    telemetry_manager = StrandsTelemetryManager()
+    telemetry_manager.setup(conversation.id, user.id)
+
     # Create ToolResultCapture to capture tool execution data
     tool_capture = ToolResultCapture(
         on_thinking=on_thinking,
@@ -104,8 +139,6 @@ def chat_with_strands(
 
     agent.callback_handler = create_callback_handler(
         on_stream=on_stream,
-        on_thinking=on_thinking,
-        on_tool_result=on_tool_result,
         on_reasoning=on_reasoning,
     )
 
@@ -186,6 +219,7 @@ def chat_with_strands(
         user=user,
         model_name=chat_input.message.model,
         continue_generate=continue_generate,
+        telemetry_manager=telemetry_manager,
         tool_capture=tool_capture,
         on_stop=on_stop,
     )

backend/app/strands_integration/handlers/callback_handler.py

@@ -39,16 +39,6 @@ def __call__(self, **kwargs):
             reasoning_text = kwargs.get("reasoningText", "")
             self.on_reasoning(reasoning_text)
             self.collected_reasoning.append(reasoning_text)
-        elif "thinking" in kwargs and self.on_reasoning:
-            thinking_text = kwargs.get("thinking", "")
-            self.on_reasoning(thinking_text)
-            self.collected_reasoning.append(thinking_text)
-        # elif "event" in kwargs:
-        #     event = kwargs["event"]
-        #     print(f"[STRANDS_CALLBACK] Event: {event}")
-        # elif "message" in kwargs:
-        #     message = kwargs["message"]
-        #     print(f"[STRANDS_CALLBACK] Message: {message}")
 
 
 def create_callback_handler(

backend/app/strands_integration/processors/result_processor.py

@@ -21,6 +21,7 @@
 
 from ..converters.message_converter import convert_strands_message_to_message_model
 from ..handlers.tool_result_capture import ToolResultCapture
+from ..telemetry.telemetry_manager import StrandsTelemetryManager
 from .cost_calculator import calculate_conversation_cost
 from .document_extractor import (
     build_thinking_log_from_tool_capture,
@@ -59,6 +60,7 @@ def post_process_strands_result(
     user: User,
     model_name: type_model_name,
     continue_generate: bool,
+    telemetry_manager: StrandsTelemetryManager,
     tool_capture: ToolResultCapture,
     on_stop: Callable[[OnStopInput], None] | None = None,
 ) -> tuple[ConversationModel, MessageModel]:
@@ -78,12 +80,21 @@ def post_process_strands_result(
     conversation.total_price += price
     conversation.should_continue = result.stop_reason == "max_tokens"
 
-    # 3. Build thinking_log from tool capture
+    # Extract reasoning content from telemetry
+    from ..telemetry import TelemetryDataExtractor
+
+    data_extractor = TelemetryDataExtractor(telemetry_manager.reasoning_processor)
+
+    reasoning_contents = data_extractor.extract_reasoning_content()
+    if reasoning_contents:
+        message.content.extend(reasoning_contents)
+
+    # Build thinking_log from tool capture
     thinking_log = build_thinking_log_from_tool_capture(tool_capture)
     if thinking_log:
         message.thinking_log = thinking_log
 
-    # 4. Set message parent and generate assistant message ID
+    # 5. Set message parent and generate assistant message ID
     message.parent = user_msg_id
 
     if continue_generate:
@@ -108,12 +119,12 @@ def post_process_strands_result(
         conversation.message_map[user_msg_id].children.append(assistant_msg_id)
         conversation.last_message_id = assistant_msg_id
 
-    # 5. Extract related documents from tool capture
+    # Extract related documents from tool capture
     related_documents = extract_related_documents_from_tool_capture(
         tool_capture, assistant_msg_id
     )
 
-    # 6. Store conversation and related documents
+    # 7. Store conversation and related documents
     store_conversation(user.id, conversation)
     if related_documents:
         store_related_documents(
@@ -122,12 +133,12 @@ def post_process_strands_result(
             related_documents=related_documents,
         )
 
-    # 7. Call on_stop callback
+    # 8. Call on_stop callback
     if on_stop:
         on_stop_input = create_on_stop_input(result, message, price)
         on_stop(on_stop_input)
 
-    # 8. Update bot statistics
+    # 9. Update bot statistics
     if bot:
         logger.debug("Bot is provided. Updating bot last used time.")
         modify_bot_last_used_time(user, bot)

backend/app/strands_integration/telemetry/__init__.py

@@ -0,0 +1,9 @@
+from .telemetry_manager import StrandsTelemetryManager
+from .processors import ReasoningSpanProcessor
+from .data_extractor import TelemetryDataExtractor
+
+__all__ = [
+    "StrandsTelemetryManager",
+    "ReasoningSpanProcessor",
+    "TelemetryDataExtractor",
+]

backend/app/strands_integration/telemetry/data_extractor.py

@@ -0,0 +1,22 @@
+"""
+Data extraction utilities for Strands telemetry.
+"""
+
+import logging
+
+from app.repositories.models.conversation import ReasoningContentModel
+
+from .processors import ReasoningSpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class TelemetryDataExtractor:
+    """Extracts structured data from telemetry span processors."""
+
+    def __init__(self, reasoning_processor: ReasoningSpanProcessor):
+        self.reasoning_processor = reasoning_processor
+
+    def extract_reasoning_content(self) -> list[ReasoningContentModel]:
+        """Extract reasoning content from telemetry data."""
+        return self.reasoning_processor.get_reasoning_data()

backend/app/strands_integration/telemetry/processors/__init__.py

@@ -0,0 +1,5 @@
+from .reasoning_processor import ReasoningSpanProcessor
+
+__all__ = [
+    "ReasoningSpanProcessor",
+]

backend/app/strands_integration/telemetry/processors/reasoning_processor.py

@@ -0,0 +1,136 @@
+"""
+Reasoning span processor for Strands telemetry.
+"""
+
+import json
+import logging
+from typing import Any, Optional
+
+from app.repositories.models.conversation import ReasoningContentModel
+from opentelemetry.context import Context
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class ReasoningSpanProcessor(SpanProcessor):
+    """Processes spans to extract reasoning content for DynamoDB storage."""
+
+    def __init__(self) -> None:
+        self.reasoning_data: list[ReasoningContentModel] = []
+        self.conversation_id: str = ""
+        self.user_id: str = ""
+
+    def set_context(self, conversation_id: str, user_id: str) -> None:
+        """Set conversation context for this processor."""
+        self.conversation_id = conversation_id
+        self.user_id = user_id
+
+    def on_start(
+        self, span: ReadableSpan, parent_context: Optional[Context] = None
+    ) -> None:
+        """Called when a span starts."""
+        pass
+
+    def on_end(self, span: ReadableSpan) -> None:
+        """Called when a span ends - extract reasoning content."""
+        if span.name == "execute_event_loop_cycle":
+            logger.debug(f"Processing Cycle span: {span.name}")
+            reasoning = self._extract_reasoning_from_span(span)
+            if reasoning:
+                self.reasoning_data.append(reasoning)
+                logger.debug(f"Extracted reasoning content from span: {span.name}")
+            else:
+                logger.debug(f"No reasoning content found in span: {span.name}")
+
+    def shutdown(self) -> None:
+        """Called when the processor is shutdown."""
+        pass
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Force flush any pending data."""
+        return True
+
+    def get_reasoning_data(self) -> list[ReasoningContentModel]:
+        """Get extracted reasoning data."""
+        return self.reasoning_data.copy()
+
+    def _extract_reasoning_from_span(
+        self, span: ReadableSpan
+    ) -> Optional[ReasoningContentModel]:
+        """
+        Extract reasoning content from span events.
+
+        Expected Data Structure:
+
+        span.events contains gen_ai.choice events with the following structure:
+
+        event.attributes["message"] = JSON string containing:
+        [
+            {
+                "reasoningContent": {
+                    "reasoningText": {
+                        "text": "The user has provided what appears to be...",
+                        "signature": "ErcBCkgIBxABGAIiQLG2dqOt..."
+                    }
+                }
+            },
+            {
+                "text": "I'll calculate the result for you."
+            },
+            {
+                "toolUse": {
+                    "toolUseId": "tooluse_xxx",
+                    "name": "calculator",
+                    "input": {"expression": "5432/64526234"}
+                }
+            }
+        ]
+        """
+        if not span.events:
+            logger.debug("No events found in span")
+            return None
+
+        for event in span.events:
+            if event.name == "gen_ai.choice":
+                if event.attributes is None:
+                    continue
+
+                logger.debug(f"Found gen_ai.choice event: {event.attributes.keys()}")
+                try:
+                    message_attr = event.attributes.get("message")
+                    if not isinstance(message_attr, str):
+                        continue
+
+                    message_content = json.loads(message_attr)
+                    logger.debug(
+                        f"Parsed message content: {len(message_content)} items"
+                    )
+
+                    for content_block in message_content:
+                        if "reasoningContent" in content_block:
+                            reasoning_data = content_block["reasoningContent"]
+                            logger.debug(
+                                f"Found reasoningContent: {reasoning_data.keys()}"
+                            )
+
+                            if "reasoningText" in reasoning_data:
+                                reasoning_text_data = reasoning_data["reasoningText"]
+                                text = reasoning_text_data.get("text", "")
+                                signature = reasoning_text_data.get("signature", "")
+
+                                if text:
+                                    logger.debug(
+                                        f"Extracted reasoning text: {len(text)} chars"
+                                    )
+                                    return ReasoningContentModel(
+                                        content_type="reasoning",
+                                        text=text,
+                                        signature=signature,
+                                        redacted_content=b"",
+                                    )
+                except (json.JSONDecodeError, KeyError) as e:
+                    logger.warning(f"Failed to parse reasoning content from event: {e}")
+
+        logger.debug("No reasoning content found in any events")
+        return None

backend/app/strands_integration/telemetry/telemetry_manager.py

@@ -0,0 +1,32 @@
+"""
+Telemetry manager for Strands integration.
+"""
+
+import logging
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from strands.telemetry import StrandsTelemetry
+from .processors import ReasoningSpanProcessor
+
+logger = logging.getLogger(__name__)
+
+
+class StrandsTelemetryManager:
+    """Manages Strands telemetry setup and span processors."""
+
+    def __init__(self):
+        self.telemetry = StrandsTelemetry()
+        self.reasoning_processor = ReasoningSpanProcessor()
+
+    def setup(self, conversation_id: str, user_id: str):
+        """Setup telemetry with custom span processors."""
+        # Setup console exporter for development
+        self.telemetry.setup_console_exporter()
+
+        # Get the tracer provider and add our custom processors
+        tracer_provider = trace.get_tracer_provider()
+        if isinstance(tracer_provider, TracerProvider):
+            tracer_provider.add_span_processor(self.reasoning_processor)
+            logger.debug("Added custom span processors to tracer provider")
+
+        self.reasoning_processor.set_context(conversation_id, user_id)

backend/tests/test_usecases/test_chat.py

@@ -907,6 +907,11 @@ class TestAgentChat(unittest.TestCase):
     model: type_model_name = "claude-v3.7-sonnet"
 
     def setUp(self) -> None:
+        # Enable debug logging for telemetry processors
+        import logging
+
+        logging.getLogger("app.strands_integration.telemetry").setLevel(logging.DEBUG)
+        logging.basicConfig(level=logging.DEBUG)
         private_bot = create_test_private_bot(
             self.bot_id,
             True,