alita-sdk 0.3.462__py3-none-any.whl → 0.3.627__py3-none-any.whl

alita_sdk/runtime/tools/loop.py

@@ -102,7 +102,9 @@ Input Data:
         logger.debug(f"LoopNode input: {predict_input}")
         completion = self.client.invoke(predict_input, config=config)
         logger.debug(f"LoopNode pure output: {completion}")
-        loop_data = _old_extract_json(completion.content.strip())
+        from ..langchain.utils import extract_text_from_completion
+        content_text = extract_text_from_completion(completion)
+        loop_data = _old_extract_json(content_text.strip())
         logger.debug(f"LoopNode output: {loop_data}")
         if self.return_type == "str":
             accumulated_response = ''

alita_sdk/runtime/tools/loop_output.py

@@ -93,7 +93,9 @@ Answer must be JSON only extractable by JSON.LOADS."""
         else:
             input_[-1].content += self.unstructured_output
             completion = self.client.invoke(input_, config=config)
-            result = _extract_json(completion.content.strip())
+            from ..langchain.utils import extract_text_from_completion
+            content_text = extract_text_from_completion(completion)
+            result = _extract_json(content_text.strip())
         try:
             tool_result: dict | List[dict] = self.tool.invoke(result, config=config, kwargs=kwargs)
             dispatch_custom_event(

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__()
@@ -85,7 +99,6 @@ class McpRemoteTool(McpServerTool):
 
     async def _execute_remote_tool(self, kwargs: Dict[str, Any]) -> str:
         """Execute the actual remote MCP tool call using SSE client."""
-        from ...tools.utils import TOOLKIT_SPLITTER
         
         # Check for session_id requirement
         if not self.session_id:
@@ -95,10 +108,10 @@ class McpRemoteTool(McpServerTool):
         # Use the original tool name from discovery for MCP server invocation
         tool_name_for_server = self.original_tool_name
         if not tool_name_for_server:
-            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}")
+            tool_name_for_server = self.name
+            logger.warning(f"original_tool_name not set for '{self.name}', using: {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 +119,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 +159,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/mcp_server_tool.py

@@ -1,11 +1,12 @@
 import uuid
 from logging import getLogger
-from typing import Any, Type, Literal, Optional, Union, List
+from typing import Any, Type, Literal, Optional, Union, List, Annotated
 
 from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field, create_model, EmailStr, constr, ConfigDict
+from pydantic import BaseModel, Field, create_model, ConfigDict, StringConstraints
 
-from ...tools.utils import TOOLKIT_SPLITTER
+# EmailStr moved to pydantic_extra_types in pydantic v2, use str for simplicity
+EmailStr = str
 
 logger = getLogger(__name__)
 
@@ -61,7 +62,7 @@ class McpServerTool(BaseTool):
                 if field.get("format") == "email":
                     return EmailStr
                 if "pattern" in field:
-                    return constr(regex=field["pattern"])
+                    return Annotated[str, StringConstraints(pattern=field["pattern"])]
                 return str
             if t == "integer":
                 return int
@@ -91,13 +92,13 @@ class McpServerTool(BaseTool):
         return create_model(model_name, **fields)
 
     def _run(self, *args, **kwargs):
-        # Extract the actual tool/prompt name (remove toolkit prefix)
+        # Use the tool name directly (no prefix extraction needed)
         call_data = {
             "server": self.server,
             "tool_timeout_sec": self.tool_timeout_sec,
             "tool_call_id": str(uuid.uuid4()),
             "params": {
-                "name": self.name.rsplit(TOOLKIT_SPLITTER)[1] if TOOLKIT_SPLITTER in self.name else self.name,
+                "name": self.name,
                 "arguments": kwargs
             }
         }

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