alita-sdk 0.3.457__py3-none-any.whl → 0.3.486__py3-none-any.whl

alita_sdk/runtime/tools/llm.py

@@ -1,3 +1,4 @@
+import asyncio
 import logging
 from traceback import format_exc
 from typing import Any, Optional, List, Union
@@ -12,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"""
     
@@ -60,6 +62,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],
@@ -132,7 +175,9 @@ class LLMNode(BaseTool):
                 struct_model = create_pydantic_model(f"LLMOutput", struct_params)
                 completion = llm_client.invoke(messages, config=config)
                 if hasattr(completion, 'tool_calls') and completion.tool_calls:
-                    new_messages, _ = self.__perform_tool_calling(completion, messages, llm_client, config)
+                    new_messages, _ = self._run_async_in_sync_context(
+                        self.__perform_tool_calling(completion, messages, llm_client, config)
+                    )
                     llm = self.__get_struct_output_model(llm_client, struct_model)
                     completion = llm.invoke(new_messages, config=config)
                     result = completion.model_dump()
@@ -155,7 +200,9 @@ class LLMNode(BaseTool):
                 # Handle both tool-calling and regular responses
                 if hasattr(completion, 'tool_calls') and completion.tool_calls:
                     # Handle iterative tool-calling and execution
-                    new_messages, current_completion = self.__perform_tool_calling(completion, messages, llm_client, config)
+                    new_messages, current_completion = self._run_async_in_sync_context(
+                        self.__perform_tool_calling(completion, messages, llm_client, config)
+                    )
 
                     output_msgs = {"messages": new_messages}
                     if self.output_variables:
@@ -190,9 +237,53 @@ class LLMNode(BaseTool):
     def _run(self, *args, **kwargs):
         # Legacy support for old interface
         return self.invoke(kwargs, **kwargs)
+    
+    def _run_async_in_sync_context(self, coro):
+        """Run async coroutine from sync context.
+        
+        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).
+        """
+        try:
+            loop = asyncio.get_running_loop()
+            # Already in async context - run in thread with new loop
+            import threading
+            
+            result_container = []
+            
+            def run_in_thread():
+                new_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(new_loop)
+                try:
+                    result_container.append(new_loop.run_until_complete(coro))
+                finally:
+                    new_loop.close()
+            
+            thread = threading.Thread(target=run_in_thread)
+            thread.start()
+            thread.join()
+            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
+            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)
+    
+    async def _arun(self, *args, **kwargs):
+        # Legacy async support
+        return self.invoke(kwargs, **kwargs)
 
-    def __perform_tool_calling(self, completion, messages, llm_client, config):
+    async def __perform_tool_calling(self, completion, messages, llm_client, config):
         # Handle iterative tool-calling and execution
+        logger.info(f"__perform_tool_calling called with {len(completion.tool_calls) if hasattr(completion, 'tool_calls') else 0} tool calls")
         new_messages = messages + [completion]
         iteration = 0
 
@@ -230,9 +321,16 @@ class LLMNode(BaseTool):
                 if tool_to_execute:
                     try:
                         logger.info(f"Executing tool '{tool_name}' with args: {tool_args}")
-                        # Pass the underlying config to the tool execution invoke method
-                        # since it may be another agent, graph, etc. to see it properly in thinking steps
-                        tool_result = tool_to_execute.invoke(tool_args, config=config)
+                        
+                        # 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")
+                            tool_result = tool_to_execute.invoke(tool_args, config=config)
 
                         # Create tool message with result - preserve structured content
                         from langchain_core.messages import ToolMessage
@@ -256,7 +354,10 @@ class LLMNode(BaseTool):
                         new_messages.append(tool_message)
 
                     except Exception as e:
-                        logger.error(f"Error executing tool '{tool_name}': {e}")
+                        import traceback
+                        error_details = traceback.format_exc()
+                        # 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(
@@ -287,16 +388,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/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",
+]

alita_sdk/runtime/tools/planning/models.py

@@ -0,0 +1,246 @@
+"""
+SQLAlchemy models for agent planning.
+
+Defines the AgentPlan table for storing execution plans with steps.
+Table is created automatically on toolkit initialization if it doesn't exist.
+"""
+
+import enum
+import logging
+import uuid
+from datetime import datetime
+from typing import List, Dict, Any, Optional
+
+from pydantic import BaseModel, Field
+from sqlalchemy import Column, String, DateTime, Text, Index, text
+from sqlalchemy.dialects.postgresql import UUID, JSONB
+from sqlalchemy.orm import declarative_base
+from sqlalchemy import create_engine
+from sqlalchemy.exc import ProgrammingError
+
+logger = logging.getLogger(__name__)
+
+Base = declarative_base()
+
+
+class PlanStatus(str, enum.Enum):
+    """Status of an execution plan."""
+    in_progress = "in_progress"
+    completed = "completed"
+    abandoned = "abandoned"
+
+
+class AgentPlan(Base):
+    """
+    Stores execution plans for agent tasks.
+    
+    Created in the project-specific pgvector database.
+    Plans are scoped by conversation_id (from server or CLI session_id).
+    """
+    __tablename__ = "agent_plans"
+    
+    id = Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4)
+    conversation_id = Column(String(255), nullable=False, index=True)
+    
+    # Plan metadata
+    title = Column(String(255), nullable=True)
+    status = Column(String(50), default=PlanStatus.in_progress.value)
+    
+    # Plan content (JSONB for flexible step storage)
+    # Structure: {"steps": [{"description": "...", "completed": false}, ...]}
+    plan_data = Column(JSONB, nullable=False, default=dict)
+    
+    # Timestamps
+    created_at = Column(DateTime, nullable=False, default=datetime.utcnow)
+    updated_at = Column(DateTime, nullable=True, onupdate=datetime.utcnow)
+
+
+# Pydantic models for tool input/output
+class PlanStep(BaseModel):
+    """A single step in a plan."""
+    description: str = Field(description="Step description")
+    completed: bool = Field(default=False, description="Whether step is completed")
+
+
+class PlanState(BaseModel):
+    """Current plan state for serialization."""
+    title: str = Field(default="", description="Plan title")
+    steps: List[PlanStep] = Field(default_factory=list, description="List of steps")
+    status: str = Field(default=PlanStatus.in_progress.value, description="Plan status")
+    
+    def render(self) -> str:
+        """Render plan as formatted string with checkboxes."""
+        if not self.steps:
+            return "No plan currently set."
+        
+        lines = []
+        if self.title:
+            lines.append(f"📋 {self.title}")
+        
+        completed_count = 0
+        for i, step in enumerate(self.steps, 1):
+            checkbox = "☑" if step.completed else "☐"
+            status_text = " (completed)" if step.completed else ""
+            lines.append(f"   {checkbox} {i}. {step.description}{status_text}")
+            if step.completed:
+                completed_count += 1
+        
+        lines.append(f"\nProgress: {completed_count}/{len(self.steps)} steps completed")
+        
+        return "\n".join(lines)
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for JSONB storage."""
+        return {
+            "steps": [{"description": s.description, "completed": s.completed} for s in self.steps]
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], title: str = "", status: str = PlanStatus.in_progress.value) -> "PlanState":
+        """Create from dictionary (JSONB data)."""
+        steps_data = data.get("steps", [])
+        steps = [PlanStep(**s) if isinstance(s, dict) else s for s in steps_data]
+        return cls(title=title, steps=steps, status=status)
+
+
+def ensure_plan_tables(connection_string: str) -> bool:
+    """
+    Ensure the agent_plans table exists in the database.
+    
+    Creates the table if it doesn't exist. Safe to call multiple times.
+    
+    Args:
+        connection_string: PostgreSQL connection string
+        
+    Returns:
+        True if table was created or already exists, False on error
+    """
+    try:
+        # Handle SecretStr if passed
+        if hasattr(connection_string, 'get_secret_value'):
+            connection_string = connection_string.get_secret_value()
+        
+        if not connection_string:
+            logger.warning("No connection string provided for plan tables")
+            return False
+            
+        engine = create_engine(connection_string)
+        
+        # Create tables if they don't exist
+        Base.metadata.create_all(engine, checkfirst=True)
+        
+        logger.debug("Agent plans table ensured")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to ensure plan tables: {e}")
+        return False
+
+
+def delete_plan_by_conversation_id(connection_string: str, conversation_id: str) -> bool:
+    """
+    Delete a plan by conversation_id.
+    
+    Args:
+        connection_string: PostgreSQL connection string
+        conversation_id: The conversation ID to delete plans for
+        
+    Returns:
+        True if deletion successful, False otherwise
+    """
+    try:
+        if hasattr(connection_string, 'get_secret_value'):
+            connection_string = connection_string.get_secret_value()
+            
+        if not connection_string or not conversation_id:
+            return False
+            
+        engine = create_engine(connection_string)
+        
+        with engine.connect() as conn:
+            result = conn.execute(
+                text("DELETE FROM agent_plans WHERE conversation_id = :conversation_id"),
+                {"conversation_id": conversation_id}
+            )
+            conn.commit()
+            
+        logger.debug(f"Deleted plan for conversation_id: {conversation_id}")
+        return True
+        
+    except Exception as e:
+        logger.error(f"Failed to delete plan for conversation_id {conversation_id}: {e}")
+        return False
+
+
+def cleanup_on_graceful_completion(
+    connection_string: str, 
+    conversation_id: str,
+    thread_id: str = None,
+    delete_checkpoints: bool = True
+) -> dict:
+    """
+    Cleanup plans and optionally checkpoints after graceful agent completion.
+    
+    This function is designed to be called after an agent completes successfully
+    (no exceptions, valid finish reason).
+    
+    Args:
+        connection_string: PostgreSQL connection string
+        conversation_id: The conversation ID to cleanup plans for
+        thread_id: The thread ID to cleanup checkpoints for (optional)
+        delete_checkpoints: If True, also delete checkpoint data
+        
+    Returns:
+        Dict with cleanup results: {'plan_deleted': bool, 'checkpoints_deleted': bool}
+    """
+    result = {'plan_deleted': False, 'checkpoints_deleted': False}
+    
+    try:
+        if hasattr(connection_string, 'get_secret_value'):
+            connection_string = connection_string.get_secret_value()
+            
+        if not connection_string or not conversation_id:
+            logger.warning("Missing connection_string or conversation_id for cleanup")
+            return result
+        
+        engine = create_engine(connection_string)
+        
+        with engine.connect() as conn:
+            # Delete plan by conversation_id
+            try:
+                conn.execute(
+                    text("DELETE FROM agent_plans WHERE conversation_id = :conversation_id"),
+                    {"conversation_id": conversation_id}
+                )
+                result['plan_deleted'] = True
+                logger.debug(f"Deleted plan for conversation_id: {conversation_id}")
+            except Exception as e:
+                # Table might not exist, which is fine
+                logger.debug(f"Could not delete plan (table may not exist): {e}")
+            
+            # Delete checkpoints if requested (still uses thread_id as that's LangGraph's key)
+            if delete_checkpoints and thread_id:
+                checkpoint_tables = [
+                    "checkpoints",
+                    "checkpoint_writes", 
+                    "checkpoint_blobs"
+                ]
+                
+                for table in checkpoint_tables:
+                    try:
+                        conn.execute(
+                            text(f"DELETE FROM {table} WHERE thread_id = :thread_id"),
+                            {"thread_id": thread_id}
+                        )
+                        logger.debug(f"Deleted {table} for thread_id: {thread_id}")
+                    except Exception as e:
+                        logger.debug(f"Could not delete from {table}: {e}")
+                
+                result['checkpoints_deleted'] = True
+            
+            conn.commit()
+            
+    except Exception as e:
+        logger.error(f"Failed to cleanup for conversation_id {conversation_id}: {e}")
+    
+    return result