lfx-nightly 0.1.13.dev0__py3-none-any.whl → 0.2.0.dev26__py3-none-any.whl

lfx/base/agents/events.py

@@ -1,4 +1,5 @@
 # Add helper functions for each event type
+import asyncio
 from collections.abc import AsyncIterator
 from time import perf_counter
 from typing import Any, Protocol
@@ -9,7 +10,7 @@ from typing_extensions import TypedDict
 
 from lfx.schema.content_block import ContentBlock
 from lfx.schema.content_types import TextContent, ToolContent
-from lfx.schema.log import SendMessageFunctionType
+from lfx.schema.log import OnTokenFunctionType, SendMessageFunctionType
 from lfx.schema.message import Message
 
 
@@ -53,7 +54,14 @@ def _calculate_duration(start_time: float) -> int:
 
 
 async def handle_on_chain_start(
-    event: dict[str, Any], agent_message: Message, send_message_method: SendMessageFunctionType, start_time: float
+    event: dict[str, Any],
+    agent_message: Message,
+    send_message_callback: SendMessageFunctionType,
+    send_token_callback: OnTokenFunctionType | None,  # noqa: ARG001
+    start_time: float,
+    *,
+    had_streaming: bool = False,  # noqa: ARG001
+    message_id: str | None = None,  # noqa: ARG001
 ) -> tuple[Message, float]:
     # Create content blocks if they don't exist
     if not agent_message.content_blocks:
@@ -80,7 +88,7 @@ async def handle_on_chain_start(
                 header={"title": "Input", "icon": "MessageSquare"},
             )
             agent_message.content_blocks[0].contents.append(text_content)
-            agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+            agent_message = await send_message_callback(message=agent_message, skip_db_update=True)
             start_time = perf_counter()
     return agent_message, start_time
 
@@ -101,15 +109,23 @@ def _extract_output_text(output: str | list) -> str:
             if isinstance(item, dict):
                 if "text" in item:
                     return item["text"] or ""
-                # If the item's type is "tool_use", return an empty string.
-                if item.get("type") == "tool_use":
-                    return ""
-                # Handle items with only 'index' key (from ChatBedrockConverse)
-                if "index" in item and len(item) == 1:
-                    return ""
-                # This is a workaround to deal with function calling by Anthropic
-                if "partial_json" in item:
+                if "content" in item:
+                    return str(item["content"])
+                if "message" in item:
+                    return str(item["message"])
+
+                # Special case handling for non-text-like dicts
+                if (
+                    item.get("type") == "tool_use"  # Handle tool use items
+                    or ("index" in item and len(item) == 1)  # Handle index-only items
+                    or "partial_json" in item  # Handle partial json items
+                    # Handle index-only items
+                    or ("index" in item and not any(k in item for k in ("text", "content", "message")))
+                    # Handle other metadata-only chunks that don't contain meaningful text
+                    or not any(key in item for key in ["text", "content", "message"])
+                ):
                     return ""
+
                 # For any other dict format, return empty string
                 return ""
             # For any other single item type (not str or dict), return empty string
@@ -133,7 +149,14 @@ def _extract_output_text(output: str | list) -> str:
 
 
 async def handle_on_chain_end(
-    event: dict[str, Any], agent_message: Message, send_message_method: SendMessageFunctionType, start_time: float
+    event: dict[str, Any],
+    agent_message: Message,
+    send_message_callback: SendMessageFunctionType,
+    send_token_callback: OnTokenFunctionType | None,  # noqa: ARG001
+    start_time: float,
+    *,
+    had_streaming: bool = False,
+    message_id: str | None = None,  # noqa: ARG001
 ) -> tuple[Message, float]:
     data_output = event["data"].get("output")
     if data_output and isinstance(data_output, AgentFinish) and data_output.return_values.get("output"):
@@ -151,7 +174,11 @@ async def handle_on_chain_end(
                 header={"title": "Output", "icon": "MessageSquare"},
             )
             agent_message.content_blocks[0].contents.append(text_content)
-        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+
+        # Only send final message if we didn't have streaming chunks
+        # If we had streaming, frontend already accumulated the chunks
+        if not had_streaming:
+            agent_message = await send_message_callback(message=agent_message)
         start_time = perf_counter()
     return agent_message, start_time
 
@@ -160,7 +187,7 @@ async def handle_on_tool_start(
     event: dict[str, Any],
     agent_message: Message,
     tool_blocks_map: dict[str, ToolContent],
-    send_message_method: SendMessageFunctionType,
+    send_message_callback: SendMessageFunctionType,
     start_time: float,
 ) -> tuple[Message, float]:
     tool_name = event["name"]
@@ -190,7 +217,7 @@ async def handle_on_tool_start(
     tool_blocks_map[tool_key] = tool_content
     agent_message.content_blocks[0].contents.append(tool_content)
 
-    agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+    agent_message = await send_message_callback(message=agent_message, skip_db_update=True)
     if agent_message.content_blocks and agent_message.content_blocks[0].contents:
         tool_blocks_map[tool_key] = agent_message.content_blocks[0].contents[-1]
     return agent_message, new_start_time
@@ -200,7 +227,7 @@ async def handle_on_tool_end(
     event: dict[str, Any],
     agent_message: Message,
     tool_blocks_map: dict[str, ToolContent],
-    send_message_method: SendMessageFunctionType,
+    send_message_callback: SendMessageFunctionType,
     start_time: float,
 ) -> tuple[Message, float]:
     run_id = event.get("run_id", "")
@@ -209,8 +236,8 @@ async def handle_on_tool_end(
     tool_content = tool_blocks_map.get(tool_key)
 
     if tool_content and isinstance(tool_content, ToolContent):
-        # Call send_message_method first to get the updated message structure
-        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+        # Call send_message_callback first to get the updated message structure
+        agent_message = await send_message_callback(message=agent_message, skip_db_update=True)
         new_start_time = perf_counter()
 
         # Now find and update the tool content in the current message
@@ -246,7 +273,7 @@ async def handle_on_tool_error(
     event: dict[str, Any],
     agent_message: Message,
     tool_blocks_map: dict[str, ToolContent],
-    send_message_method: SendMessageFunctionType,
+    send_message_callback: SendMessageFunctionType,
     start_time: float,
 ) -> tuple[Message, float]:
     run_id = event.get("run_id", "")
@@ -258,7 +285,7 @@ async def handle_on_tool_error(
         tool_content.error = event["data"].get("error", "Unknown error")
         tool_content.duration = _calculate_duration(start_time)
         tool_content.header = {"title": f"Error using **{tool_content.name}**", "icon": "Hammer"}
-        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+        agent_message = await send_message_callback(message=agent_message, skip_db_update=True)
         start_time = perf_counter()
     return agent_message, start_time
 
@@ -266,8 +293,12 @@ async def handle_on_tool_error(
 async def handle_on_chain_stream(
     event: dict[str, Any],
     agent_message: Message,
-    send_message_method: SendMessageFunctionType,
+    send_message_callback: SendMessageFunctionType,  # noqa: ARG001
+    send_token_callback: OnTokenFunctionType | None,
     start_time: float,
+    *,
+    had_streaming: bool = False,  # noqa: ARG001
+    message_id: str | None = None,
 ) -> tuple[Message, float]:
     data_chunk = event["data"].get("chunk", {})
     if isinstance(data_chunk, dict) and data_chunk.get("output"):
@@ -275,15 +306,26 @@ async def handle_on_chain_stream(
         if output and isinstance(output, str | list):
             agent_message.text = _extract_output_text(output)
         agent_message.properties.state = "complete"
-        agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+        # Don't call send_message_callback here - we must update in place
+        # in order to keep the message id consistent throughout the stream.
+        # The final message will be sent after the loop completes
         start_time = perf_counter()
     elif isinstance(data_chunk, AIMessageChunk):
         output_text = _extract_output_text(data_chunk.content)
-        if output_text and isinstance(agent_message.text, str):
-            agent_message.text += output_text
-            agent_message.properties.state = "partial"
-            agent_message = await send_message_method(message=agent_message, skip_db_update=True)
+        # For streaming, send token event if callback is available
+        # Note: we should expect the callback, but we keep it optional for backwards compatibility
+        # as of v1.6.5
+        if output_text and output_text.strip() and send_token_callback and message_id:
+            await asyncio.to_thread(
+                send_token_callback,
+                data={
+                    "chunk": output_text,
+                    "id": str(message_id),
+                },
+            )
+
         if not agent_message.text:
+            # Starts the timer when the first message is starting to be generated
             start_time = perf_counter()
     return agent_message, start_time
 
@@ -294,7 +336,7 @@ class ToolEventHandler(Protocol):
         event: dict[str, Any],
         agent_message: Message,
         tool_blocks_map: dict[str, ContentBlock],
-        send_message_method: SendMessageFunctionType,
+        send_message_callback: SendMessageFunctionType,
         start_time: float,
     ) -> tuple[Message, float]: ...
 
@@ -304,8 +346,12 @@ class ChainEventHandler(Protocol):
         self,
         event: dict[str, Any],
         agent_message: Message,
-        send_message_method: SendMessageFunctionType,
+        send_message_callback: SendMessageFunctionType,
+        send_token_callback: OnTokenFunctionType | None,
         start_time: float,
+        *,
+        had_streaming: bool = False,
+        message_id: str | None = None,
     ) -> tuple[Message, float]: ...
 
 
@@ -329,7 +375,8 @@ TOOL_EVENT_HANDLERS: dict[str, ToolEventHandler] = {
 async def process_agent_events(
     agent_executor: AsyncIterator[dict[str, Any]],
     agent_message: Message,
-    send_message_method: SendMessageFunctionType,
+    send_message_callback: SendMessageFunctionType,
+    send_token_callback: OnTokenFunctionType | None = None,
 ) -> Message:
     """Process agent events and return the final output."""
     if isinstance(agent_message.properties, dict):
@@ -337,26 +384,47 @@ async def process_agent_events(
     else:
         agent_message.properties.icon = "Bot"
         agent_message.properties.state = "partial"
-    # Store the initial message
-    agent_message = await send_message_method(message=agent_message)
+    # Store the initial message and capture the message id
+    agent_message = await send_message_callback(message=agent_message)
+    # Capture the original message id - this must stay consistent throughout if streaming
+    # Message may not contain id if the Agent is not connected to a Chat Output (_should_skip_message is True)
+    initial_message_id = agent_message.id if hasattr(agent_message, "id") else None
     try:
         # Create a mapping of run_ids to tool contents
         tool_blocks_map: dict[str, ToolContent] = {}
+        had_streaming = False
         start_time = perf_counter()
+
         async for event in agent_executor:
             if event["event"] in TOOL_EVENT_HANDLERS:
                 tool_handler = TOOL_EVENT_HANDLERS[event["event"]]
                 # Use skip_db_update=True during streaming to avoid DB round-trips
                 agent_message, start_time = await tool_handler(
-                    event, agent_message, tool_blocks_map, send_message_method, start_time
+                    event, agent_message, tool_blocks_map, send_message_callback, start_time
                 )
             elif event["event"] in CHAIN_EVENT_HANDLERS:
                 chain_handler = CHAIN_EVENT_HANDLERS[event["event"]]
-                # Use skip_db_update=True during streaming to avoid DB round-trips
-                agent_message, start_time = await chain_handler(event, agent_message, send_message_method, start_time)
+
+                # Check if this is a streaming event
+                if event["event"] in ("on_chain_stream", "on_chat_model_stream"):
+                    had_streaming = True
+                    agent_message, start_time = await chain_handler(
+                        event,
+                        agent_message,
+                        send_message_callback,
+                        send_token_callback,
+                        start_time,
+                        had_streaming=had_streaming,
+                        message_id=initial_message_id,
+                    )
+                else:
+                    agent_message, start_time = await chain_handler(
+                        event, agent_message, send_message_callback, None, start_time, had_streaming=had_streaming
+                    )
+
         agent_message.properties.state = "complete"
         # Final DB update with the complete message (skip_db_update=False by default)
-        agent_message = await send_message_method(message=agent_message)
+        agent_message = await send_message_callback(message=agent_message)
     except Exception as e:
         raise ExceptionWithMessageError(agent_message, str(e)) from e
     return await Message.create(**agent_message.model_dump())

lfx/base/agents/utils.py

@@ -47,9 +47,22 @@ def data_to_messages(data: list[Data | Message]) -> list[BaseMessage]:
         data (List[Data | Message]): The data to convert.
 
     Returns:
-        List[BaseMessage]: The data as messages.
+        List[BaseMessage]: The data as messages, filtering out any with empty content.
     """
-    return [value.to_lc_message() for value in data]
+    messages = []
+    for value in data:
+        try:
+            lc_message = value.to_lc_message()
+            # Only add messages with non-empty content (prevents Anthropic API errors)
+            content = lc_message.content
+            if content and ((isinstance(content, str) and content.strip()) or (isinstance(content, list) and content)):
+                messages.append(lc_message)
+            else:
+                logger.warning("Skipping message with empty content in chat history")
+        except (ValueError, AttributeError) as e:
+            logger.warning(f"Failed to convert message to BaseMessage: {e}")
+            continue
+    return messages
 
 
 def validate_and_create_xml_agent(