alita-sdk 0.3.465__py3-none-any.whl → 0.3.497__py3-none-any.whl

alita_sdk/runtime/tools/llm.py

@@ -13,6 +13,7 @@ from ..langchain.utils import create_pydantic_model, propagate_the_input_mapping
 
 logger = logging.getLogger(__name__)
 
+
 class LLMNode(BaseTool):
     """Enhanced LLM node with chat history and tool binding support"""
     
@@ -33,6 +34,7 @@ class LLMNode(BaseTool):
     available_tools: Optional[List[BaseTool]] = Field(default=None, description='Available tools for binding')
     tool_names: Optional[List[str]] = Field(default=None, description='Specific tool names to filter')
     steps_limit: Optional[int] = Field(default=25, description='Maximum steps for tool execution')
+    tool_execution_timeout: Optional[int] = Field(default=900, description='Timeout (seconds) for tool execution. Default is 15 minutes.')
 
     def get_filtered_tools(self) -> List[BaseTool]:
         """
@@ -61,6 +63,47 @@ class LLMNode(BaseTool):
         
         return filtered_tools
 
+    def _get_tool_truncation_suggestions(self, tool_name: Optional[str]) -> str:
+        """
+        Get context-specific suggestions for how to reduce output from a tool.
+        
+        First checks if the tool itself provides truncation suggestions via 
+        `truncation_suggestions` attribute or `get_truncation_suggestions()` method.
+        Falls back to generic suggestions if the tool doesn't provide any.
+        
+        Args:
+            tool_name: Name of the tool that caused the context overflow
+            
+        Returns:
+            Formatted string with numbered suggestions for the specific tool
+        """
+        suggestions = None
+        
+        # Try to get suggestions from the tool itself
+        if tool_name:
+            filtered_tools = self.get_filtered_tools()
+            for tool in filtered_tools:
+                if tool.name == tool_name:
+                    # Check for truncation_suggestions attribute
+                    if hasattr(tool, 'truncation_suggestions') and tool.truncation_suggestions:
+                        suggestions = tool.truncation_suggestions
+                        break
+                    # Check for get_truncation_suggestions method
+                    elif hasattr(tool, 'get_truncation_suggestions') and callable(tool.get_truncation_suggestions):
+                        suggestions = tool.get_truncation_suggestions()
+                        break
+        
+        # Fall back to generic suggestions if tool doesn't provide any
+        if not suggestions:
+            suggestions = [
+                "Check if the tool has parameters to limit output size (e.g., max_items, max_results, max_depth)",
+                "Target a more specific path or query instead of broad searches",
+                "Break the operation into smaller, focused requests",
+            ]
+        
+        # Format as numbered list
+        return "\n".join(f"{i+1}. {s}" for i, s in enumerate(suggestions))
+
     def invoke(
             self,
             state: Union[str, dict],
@@ -87,7 +130,9 @@ class LLMNode(BaseTool):
         # or standalone LLM node for chat (with messages only)
         if 'system' in func_args.keys():
             # Flow for LLM node with prompt/task from pipeline
-            if not func_args.get('system') or not func_args.get('task'):
+            if func_args.get('system') is None or func_args.get('task') is None:
+                raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
+                                    f"Actual params: {func_args}")
                 raise ToolException(f"LLMNode requires 'system' and 'task' parameters in input mapping. "
                                     f"Actual params: {func_args}")
             # cast to str in case user passes variable different from str
@@ -201,40 +246,146 @@ class LLMNode(BaseTool):
         
         For MCP tools with persistent sessions, we reuse the same event loop
         that was used to create the MCP client and sessions (set by CLI).
+
+        When called from within a running event loop (e.g., nested LLM nodes),
+        we need to handle this carefully to avoid "event loop already running" errors.
+
+        This method handles three scenarios:
+        1. Called from async context (event loop running) - creates new thread with new loop
+        2. Called from sync context with persistent loop - reuses persistent loop
+        3. Called from sync context without loop - creates new persistent loop
         """
+        import threading
+
+        # Check if there's a running loop
         try:
-            loop = asyncio.get_running_loop()
-            # Already in async context - run in thread with new loop
-            import threading
-            
+            running_loop = asyncio.get_running_loop()
+            loop_is_running = True
+            logger.debug(f"Detected running event loop (id: {id(running_loop)}), executing tool calls in separate thread")
+        except RuntimeError:
+            loop_is_running = False
+
+        # Scenario 1: Loop is currently running - MUST use thread
+        if loop_is_running:
             result_container = []
-            
+            exception_container = []
+
+            # Try to capture Streamlit context from current thread for propagation
+            streamlit_ctx = None
+            try:
+                from streamlit.runtime.scriptrunner import get_script_run_ctx, add_script_run_ctx
+                streamlit_ctx = get_script_run_ctx()
+                if streamlit_ctx:
+                    logger.debug("Captured Streamlit context for propagation to worker thread")
+            except (ImportError, Exception) as e:
+                logger.debug(f"Streamlit context not available or failed to capture: {e}")
+
             def run_in_thread():
+                """Run coroutine in a new thread with its own event loop."""
                 new_loop = asyncio.new_event_loop()
                 asyncio.set_event_loop(new_loop)
                 try:
-                    result_container.append(new_loop.run_until_complete(coro))
+                    result = new_loop.run_until_complete(coro)
+                    result_container.append(result)
+                except Exception as e:
+                    logger.debug(f"Exception in async thread: {e}")
+                    exception_container.append(e)
                 finally:
                     new_loop.close()
-            
-            thread = threading.Thread(target=run_in_thread)
+                    asyncio.set_event_loop(None)
+
+            thread = threading.Thread(target=run_in_thread, daemon=False)
+
+            # Propagate Streamlit context to the worker thread if available
+            if streamlit_ctx is not None:
+                try:
+                    add_script_run_ctx(thread, streamlit_ctx)
+                    logger.debug("Successfully propagated Streamlit context to worker thread")
+                except Exception as e:
+                    logger.warning(f"Failed to propagate Streamlit context to worker thread: {e}")
+
             thread.start()
-            thread.join()
+            thread.join(timeout=self.tool_execution_timeout)  # 15 minute timeout for safety
+
+            if thread.is_alive():
+                logger.error("Async operation timed out after 5 minutes")
+                raise TimeoutError("Async operation in thread timed out")
+
+            # Re-raise exception if one occurred
+            if exception_container:
+                raise exception_container[0]
+
             return result_container[0] if result_container else None
-            
-        except RuntimeError:
-            # No event loop running - use/create persistent loop
-            # This loop is shared with MCP session creation for stateful tools
+
+        # Scenario 2 & 3: No loop running - use or create persistent loop
+        else:
+            # Get or create persistent loop
             if not hasattr(self.__class__, '_persistent_loop') or \
                self.__class__._persistent_loop is None or \
                self.__class__._persistent_loop.is_closed():
                 self.__class__._persistent_loop = asyncio.new_event_loop()
                 logger.debug("Created persistent event loop for async tools")
-            
+
             loop = self.__class__._persistent_loop
-            asyncio.set_event_loop(loop)
-            return loop.run_until_complete(coro)
-    
+
+            # Double-check the loop is not running (safety check)
+            if loop.is_running():
+                logger.debug("Persistent loop is unexpectedly running, using thread execution")
+
+                result_container = []
+                exception_container = []
+
+                # Try to capture Streamlit context from current thread for propagation
+                streamlit_ctx = None
+                try:
+                    from streamlit.runtime.scriptrunner import get_script_run_ctx, add_script_run_ctx
+                    streamlit_ctx = get_script_run_ctx()
+                    if streamlit_ctx:
+                        logger.debug("Captured Streamlit context for propagation to worker thread")
+                except (ImportError, Exception) as e:
+                    logger.debug(f"Streamlit context not available or failed to capture: {e}")
+
+                def run_in_thread():
+                    """Run coroutine in a new thread with its own event loop."""
+                    new_loop = asyncio.new_event_loop()
+                    asyncio.set_event_loop(new_loop)
+                    try:
+                        result = new_loop.run_until_complete(coro)
+                        result_container.append(result)
+                    except Exception as ex:
+                        logger.debug(f"Exception in async thread: {ex}")
+                        exception_container.append(ex)
+                    finally:
+                        new_loop.close()
+                        asyncio.set_event_loop(None)
+
+                thread = threading.Thread(target=run_in_thread, daemon=False)
+
+                # Propagate Streamlit context to the worker thread if available
+                if streamlit_ctx is not None:
+                    try:
+                        add_script_run_ctx(thread, streamlit_ctx)
+                        logger.debug("Successfully propagated Streamlit context to worker thread")
+                    except Exception as e:
+                        logger.warning(f"Failed to propagate Streamlit context to worker thread: {e}")
+
+                thread.start()
+                thread.join(timeout=self.tool_execution_timeout)
+
+                if thread.is_alive():
+                    logger.error("Async operation timed out after 15 minutes")
+                    raise TimeoutError("Async operation in thread timed out")
+
+                if exception_container:
+                    raise exception_container[0]
+
+                return result_container[0] if result_container else None
+            else:
+                # Loop exists but not running - safe to use run_until_complete
+                logger.debug(f"Using persistent loop (id: {id(loop)}) with run_until_complete")
+                asyncio.set_event_loop(loop)
+                return loop.run_until_complete(coro)
+
     async def _arun(self, *args, **kwargs):
         # Legacy async support
         return self.invoke(kwargs, **kwargs)
@@ -282,12 +433,14 @@ class LLMNode(BaseTool):
                         
                         # Try async invoke first (for MCP tools), fallback to sync
                         tool_result = None
-                        try:
-                            # Try async invocation first
-                            tool_result = await tool_to_execute.ainvoke(tool_args, config=config)
-                        except NotImplementedError:
-                            # Tool doesn't support async, use sync invoke
-                            logger.debug(f"Tool '{tool_name}' doesn't support async, using sync invoke")
+                        if hasattr(tool_to_execute, 'ainvoke'):
+                            try:
+                                tool_result = await tool_to_execute.ainvoke(tool_args, config=config)
+                            except (NotImplementedError, AttributeError):
+                                logger.debug(f"Tool '{tool_name}' ainvoke failed, falling back to sync invoke")
+                                tool_result = tool_to_execute.invoke(tool_args, config=config)
+                        else:
+                            # Sync-only tool
                             tool_result = tool_to_execute.invoke(tool_args, config=config)
 
                         # Create tool message with result - preserve structured content
@@ -314,7 +467,8 @@ class LLMNode(BaseTool):
                     except Exception as e:
                         import traceback
                         error_details = traceback.format_exc()
-                        logger.error(f"Error executing tool '{tool_name}': {e}\n{error_details}")
+                        # Use debug level to avoid duplicate output when CLI callbacks are active
+                        logger.debug(f"Error executing tool '{tool_name}': {e}\n{error_details}")
                         # Create error tool message
                         from langchain_core.messages import ToolMessage
                         tool_message = ToolMessage(
@@ -345,16 +499,150 @@ class LLMNode(BaseTool):
                     break
 
             except Exception as e:
-                logger.error(f"Error in LLM call during iteration {iteration}: {e}")
-                # Add error message and break the loop
-                error_msg = f"Error processing tool results in iteration {iteration}: {str(e)}"
-                new_messages.append(AIMessage(content=error_msg))
-                break
+                error_str = str(e).lower()
+                
+                # Check for context window / token limit errors
+                is_context_error = any(indicator in error_str for indicator in [
+                    'context window', 'context_window', 'token limit', 'too long',
+                    'maximum context length', 'input is too long', 'exceeds the limit',
+                    'contextwindowexceedederror', 'max_tokens', 'content too large'
+                ])
+                
+                # Check for Bedrock/Claude output limit errors
+                # These often manifest as "model identifier is invalid" when output exceeds limits
+                is_output_limit_error = any(indicator in error_str for indicator in [
+                    'model identifier is invalid',
+                    'bedrockexception',
+                    'output token',
+                    'response too large',
+                    'max_tokens_to_sample',
+                    'output_token_limit'
+                ])
+                
+                if is_context_error or is_output_limit_error:
+                    error_type = "output limit" if is_output_limit_error else "context window"
+                    logger.warning(f"{error_type.title()} exceeded during tool execution iteration {iteration}")
+                    
+                    # Find the last tool message and its associated tool name
+                    last_tool_msg_idx = None
+                    last_tool_name = None
+                    last_tool_call_id = None
+                    
+                    # First, find the last tool message
+                    for i in range(len(new_messages) - 1, -1, -1):
+                        msg = new_messages[i]
+                        if hasattr(msg, 'tool_call_id') or (hasattr(msg, 'type') and getattr(msg, 'type', None) == 'tool'):
+                            last_tool_msg_idx = i
+                            last_tool_call_id = getattr(msg, 'tool_call_id', None)
+                            break
+                    
+                    # Find the tool name from the AIMessage that requested this tool call
+                    if last_tool_call_id:
+                        for i in range(last_tool_msg_idx - 1, -1, -1):
+                            msg = new_messages[i]
+                            if hasattr(msg, 'tool_calls') and msg.tool_calls:
+                                for tc in msg.tool_calls:
+                                    tc_id = tc.get('id', '') if isinstance(tc, dict) else getattr(tc, 'id', '')
+                                    if tc_id == last_tool_call_id:
+                                        last_tool_name = tc.get('name', '') if isinstance(tc, dict) else getattr(tc, 'name', '')
+                                        break
+                                if last_tool_name:
+                                    break
+                    
+                    # Build dynamic suggestion based on the tool that caused the overflow
+                    tool_suggestions = self._get_tool_truncation_suggestions(last_tool_name)
+                    
+                    # Truncate the problematic tool result if found
+                    if last_tool_msg_idx is not None:
+                        from langchain_core.messages import ToolMessage
+                        original_msg = new_messages[last_tool_msg_idx]
+                        tool_call_id = getattr(original_msg, 'tool_call_id', 'unknown')
+                        
+                        # Build error-specific guidance
+                        if is_output_limit_error:
+                            truncated_content = (
+                                f"⚠️ MODEL OUTPUT LIMIT EXCEEDED\n\n"
+                                f"The tool '{last_tool_name or 'unknown'}' returned data, but the model's response was too large.\n\n"
+                                f"IMPORTANT: You must provide a SMALLER, more focused response.\n"
+                                f"- Break down your response into smaller chunks\n"
+                                f"- Summarize instead of listing everything\n"
+                                f"- Focus on the most relevant information first\n"
+                                f"- If listing items, show only top 5-10 most important\n\n"
+                                f"Tool-specific tips:\n{tool_suggestions}\n\n"
+                                f"Please retry with a more concise response."
+                            )
+                        else:
+                            truncated_content = (
+                                f"⚠️ TOOL OUTPUT TRUNCATED - Context window exceeded\n\n"
+                                f"The tool '{last_tool_name or 'unknown'}' returned too much data for the model's context window.\n\n"
+                                f"To fix this:\n{tool_suggestions}\n\n"
+                                f"Please retry with more restrictive parameters."
+                            )
+                        
+                        truncated_msg = ToolMessage(
+                            content=truncated_content,
+                            tool_call_id=tool_call_id
+                        )
+                        new_messages[last_tool_msg_idx] = truncated_msg
+                        
+                        logger.info(f"Truncated large tool result from '{last_tool_name}' and continuing")
+                        # Continue to next iteration - the model will see the truncation message
+                        continue
+                    else:
+                        # Couldn't find tool message, add error and break
+                        if is_output_limit_error:
+                            error_msg = (
+                                "Model output limit exceeded. Please provide a more concise response. "
+                                "Break down your answer into smaller parts and summarize where possible."
+                            )
+                        else:
+                            error_msg = (
+                                "Context window exceeded. The conversation or tool results are too large. "
+                                "Try using tools with smaller output limits (e.g., max_items, max_depth parameters)."
+                            )
+                        new_messages.append(AIMessage(content=error_msg))
+                        break
+                else:
+                    logger.error(f"Error in LLM call during iteration {iteration}: {e}")
+                    # Add error message and break the loop
+                    error_msg = f"Error processing tool results in iteration {iteration}: {str(e)}"
+                    new_messages.append(AIMessage(content=error_msg))
+                    break
 
-        # Log completion status
+        # Handle max iterations
         if iteration >= self.steps_limit:
             logger.warning(f"Reached maximum iterations ({self.steps_limit}) for tool execution")
-            # Add a warning message to the chat
+            
+            # CRITICAL: Check if the last message is an AIMessage with pending tool_calls
+            # that were not processed. If so, we need to add placeholder ToolMessages to prevent
+            # the "assistant message with 'tool_calls' must be followed by tool messages" error
+            # when the conversation continues.
+            if new_messages:
+                last_msg = new_messages[-1]
+                if hasattr(last_msg, 'tool_calls') and last_msg.tool_calls:
+                    from langchain_core.messages import ToolMessage
+                    pending_tool_calls = last_msg.tool_calls if hasattr(last_msg.tool_calls, '__iter__') else []
+                    
+                    # Check which tool_call_ids already have responses
+                    existing_tool_call_ids = set()
+                    for msg in new_messages:
+                        if hasattr(msg, 'tool_call_id'):
+                            existing_tool_call_ids.add(msg.tool_call_id)
+                    
+                    # Add placeholder responses for any tool calls without responses
+                    for tool_call in pending_tool_calls:
+                        tool_call_id = tool_call.get('id', '') if isinstance(tool_call, dict) else getattr(tool_call, 'id', '')
+                        tool_name = tool_call.get('name', '') if isinstance(tool_call, dict) else getattr(tool_call, 'name', '')
+                        
+                        if tool_call_id and tool_call_id not in existing_tool_call_ids:
+                            logger.info(f"Adding placeholder ToolMessage for interrupted tool call: {tool_name} ({tool_call_id})")
+                            placeholder_msg = ToolMessage(
+                                content=f"[Tool execution interrupted - step limit ({self.steps_limit}) reached before {tool_name} could be executed]",
+                                tool_call_id=tool_call_id
+                            )
+                            new_messages.append(placeholder_msg)
+            
+            # Add warning message - CLI or calling code can detect this and prompt user
             warning_msg = f"Maximum tool execution iterations ({self.steps_limit}) reached. Stopping tool execution."
             new_messages.append(AIMessage(content=warning_msg))
         else:

alita_sdk/runtime/tools/mcp_remote_tool.py

@@ -20,10 +20,14 @@ from ..utils.mcp_oauth import (
     fetch_resource_metadata_async,
     infer_authorization_servers_from_realm,
 )
-from ..utils.mcp_sse_client import McpSseClient
+from ..utils.mcp_client import McpClient
 
 logger = logging.getLogger(__name__)
 
+# Global registry to store MCP tool session metadata by tool name
+# This is used to pass session info to callbacks since LangChain's serialization doesn't include all fields
+MCP_TOOL_SESSION_REGISTRY: Dict[str, Dict[str, Any]] = {}
+
 
 class McpRemoteTool(McpServerTool):
     """
@@ -43,6 +47,7 @@ class McpRemoteTool(McpServerTool):
         """Update metadata with session info after model initialization."""
         super().model_post_init(__context)
         self._update_metadata_with_session()
+        self._register_session_metadata()
     
     def _update_metadata_with_session(self):
         """Update the metadata dict with current session information."""
@@ -54,6 +59,15 @@ class McpRemoteTool(McpServerTool):
                 'mcp_server_url': canonical_resource(self.server_url)
             })
     
+    def _register_session_metadata(self):
+        """Register session metadata in global registry for callback access."""
+        if self.session_id and self.server_url:
+            MCP_TOOL_SESSION_REGISTRY[self.name] = {
+                'mcp_session_id': self.session_id,
+                'mcp_server_url': canonical_resource(self.server_url)
+            }
+            logger.debug(f"[MCP] Registered session metadata for tool '{self.name}': session={self.session_id}")
+    
     def __getstate__(self):
         """Custom serialization for pickle compatibility."""
         state = super().__getstate__()
@@ -98,7 +112,7 @@ class McpRemoteTool(McpServerTool):
             tool_name_for_server = self.name.rsplit(TOOLKIT_SPLITTER, 1)[-1] if TOOLKIT_SPLITTER in self.name else self.name
             logger.warning(f"original_tool_name not set for '{self.name}', using extracted: {tool_name_for_server}")
         
-        logger.info(f"[MCP SSE] Executing tool '{tool_name_for_server}' with session {self.session_id}")
+        logger.info(f"[MCP] Executing tool '{tool_name_for_server}' with session {self.session_id}")
         
         try:
             # Prepare headers
@@ -106,16 +120,18 @@ class McpRemoteTool(McpServerTool):
             if self.server_headers:
                 headers.update(self.server_headers)
             
-            # Create SSE client
-            client = McpSseClient(
+            # Create unified MCP client (auto-detects transport)
+            client = McpClient(
                 url=self.server_url,
                 session_id=self.session_id,
                 headers=headers,
                 timeout=self.tool_timeout_sec
             )
             
-            # Execute tool call via SSE
-            result = await client.call_tool(tool_name_for_server, kwargs)
+            # Execute tool call (client auto-detects SSE vs Streamable HTTP)
+            async with client:
+                await client.initialize()
+                result = await client.call_tool(tool_name_for_server, kwargs)
             
             # Format the result
             if isinstance(result, dict):
@@ -144,7 +160,7 @@ class McpRemoteTool(McpServerTool):
             return str(result)
             
         except Exception as e:
-            logger.error(f"[MCP SSE] Tool execution failed: {e}", exc_info=True)
+            logger.error(f"[MCP] Tool execution failed: {e}", exc_info=True)
             raise
 
     def _parse_sse(self, text: str) -> Dict[str, Any]:

alita_sdk/runtime/tools/planning/__init__.py

@@ -0,0 +1,36 @@
+"""
+Planning tools for runtime agents.
+
+Provides plan management for multi-step task execution with progress tracking.
+Supports two storage backends:
+1. PostgreSQL - when connection_string is provided (production/indexer_worker)
+2. Filesystem - when no connection string (local CLI usage)
+"""
+
+from .wrapper import (
+    PlanningWrapper,
+    PlanStep,
+    PlanState,
+    FilesystemStorage,
+    PostgresStorage,
+)
+from .models import (
+    AgentPlan, 
+    PlanStatus, 
+    ensure_plan_tables,
+    delete_plan_by_conversation_id,
+    cleanup_on_graceful_completion
+)
+
+__all__ = [
+    "PlanningWrapper",
+    "PlanStep",
+    "PlanState",
+    "FilesystemStorage",
+    "PostgresStorage",
+    "AgentPlan",
+    "PlanStatus",
+    "ensure_plan_tables",
+    "delete_plan_by_conversation_id",
+    "cleanup_on_graceful_completion",
+]