alita-sdk 0.3.376__py3-none-any.whl → 0.3.435__py3-none-any.whl

alita_sdk/runtime/tools/llm.py

@@ -7,6 +7,7 @@ from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool, ToolException
 from pydantic import Field
 
+from ..langchain.constants import ELITEA_RS
 from ..langchain.utils import create_pydantic_model, propagate_the_input_mapping
 
 logger = logging.getLogger(__name__)
@@ -30,6 +31,7 @@ class LLMNode(BaseTool):
     structured_output: Optional[bool] = Field(default=False, description='Whether to use structured output')
     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')
 
     def get_filtered_tools(self) -> List[BaseTool]:
         """
@@ -88,8 +90,11 @@ class LLMNode(BaseTool):
                 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
-            messages = [SystemMessage(content=str(func_args.get('system'))), HumanMessage(content=str(func_args.get('task')))]
-            messages.extend(func_args.get('chat_history', []))
+            messages = [SystemMessage(content=str(func_args.get('system'))), *func_args.get('chat_history', []), HumanMessage(content=str(func_args.get('task')))]
+            # Remove pre-last item if last two messages are same type and content
+            if len(messages) >= 2 and type(messages[-1]) == type(messages[-2]) and messages[-1].content == messages[
+                -2].content:
+                messages.pop(-2)
         else:
             # Flow for chat-based LLM node w/o prompt/task from pipeline but with messages in state
             # verify messages structure
@@ -122,14 +127,25 @@ class LLMNode(BaseTool):
                     }
                     for key, value in (self.structured_output_dict or {}).items()
                 }
+                # Add default output field for proper response to user
+                struct_params['elitea_response'] = {'description': 'final output to user', 'type': 'str'}
                 struct_model = create_pydantic_model(f"LLMOutput", struct_params)
-                llm = llm_client.with_structured_output(struct_model)
-                completion = llm.invoke(messages, config=config)
-                result = completion.model_dump()
+                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)
+                    llm = self.__get_struct_output_model(llm_client, struct_model)
+                    completion = llm.invoke(new_messages, config=config)
+                    result = completion.model_dump()
+                else:
+                    llm = self.__get_struct_output_model(llm_client, struct_model)
+                    completion = llm.invoke(messages, config=config)
+                    result = completion.model_dump()
 
                 # Ensure messages are properly formatted
                 if result.get('messages') and isinstance(result['messages'], list):
                     result['messages'] = [{'role': 'assistant', 'content': '\n'.join(result['messages'])}]
+                else:
+                    result['messages'] = messages + [AIMessage(content=result.get(ELITEA_RS, ''))]
 
                 return result
             else:
@@ -139,115 +155,15 @@ 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 = messages + [completion]
-                    max_iterations = 15
-                    iteration = 0
-
-                    # Continue executing tools until no more tool calls or max iterations reached
-                    current_completion = completion
-                    while (hasattr(current_completion, 'tool_calls') and
-                           current_completion.tool_calls and
-                           iteration < max_iterations):
-
-                        iteration += 1
-                        logger.info(f"Tool execution iteration {iteration}/{max_iterations}")
-
-                        # Execute each tool call in the current completion
-                        tool_calls = current_completion.tool_calls if hasattr(current_completion.tool_calls,
-                                                                              '__iter__') else []
-
-                        for tool_call in tool_calls:
-                            tool_name = tool_call.get('name', '') if isinstance(tool_call, dict) else getattr(tool_call,
-                                                                                                              'name',
-                                                                                                              '')
-                            tool_args = tool_call.get('args', {}) if isinstance(tool_call, dict) else getattr(tool_call,
-                                                                                                              'args',
-                                                                                                              {})
-                            tool_call_id = tool_call.get('id', '') if isinstance(tool_call, dict) else getattr(
-                                tool_call, 'id', '')
-
-                            # Find the tool in filtered tools
-                            filtered_tools = self.get_filtered_tools()
-                            tool_to_execute = None
-                            for tool in filtered_tools:
-                                if tool.name == tool_name:
-                                    tool_to_execute = tool
-                                    break
-
-                            if tool_to_execute:
-                                try:
-                                    logger.info(f"Executing tool '{tool_name}' with args: {tool_args}")
-                                    tool_result = tool_to_execute.invoke(tool_args)
-
-                                    # Create tool message with result - preserve structured content
-                                    from langchain_core.messages import ToolMessage
-                                    
-                                    # Check if tool_result is structured content (list of dicts)
-                                    # TODO: need solid check for being compatible with ToolMessage content format
-                                    if isinstance(tool_result, list) and all(
-                                        isinstance(item, dict) and 'type' in item for item in tool_result
-                                    ):
-                                        # Use structured content directly for multimodal support
-                                        tool_message = ToolMessage(
-                                            content=tool_result,
-                                            tool_call_id=tool_call_id
-                                        )
-                                    else:
-                                        # Fallback to string conversion for other tool results
-                                        tool_message = ToolMessage(
-                                            content=str(tool_result),
-                                            tool_call_id=tool_call_id
-                                        )
-                                    new_messages.append(tool_message)
-
-                                except Exception as e:
-                                    logger.error(f"Error executing tool '{tool_name}': {e}")
-                                    # Create error tool message
-                                    from langchain_core.messages import ToolMessage
-                                    tool_message = ToolMessage(
-                                        content=f"Error executing {tool_name}: {str(e)}",
-                                        tool_call_id=tool_call_id
-                                    )
-                                    new_messages.append(tool_message)
-                            else:
-                                logger.warning(f"Tool '{tool_name}' not found in available tools")
-                                # Create error tool message for missing tool
-                                from langchain_core.messages import ToolMessage
-                                tool_message = ToolMessage(
-                                    content=f"Tool '{tool_name}' not available",
-                                    tool_call_id=tool_call_id
-                                )
-                                new_messages.append(tool_message)
-
-                        # Call LLM again with tool results to get next response
-                        try:
-                            current_completion = llm_client.invoke(new_messages, config=config)
-                            new_messages.append(current_completion)
-
-                            # Check if we still have tool calls
-                            if hasattr(current_completion, 'tool_calls') and current_completion.tool_calls:
-                                logger.info(f"LLM requested {len(current_completion.tool_calls)} more tool calls")
-                            else:
-                                logger.info("LLM completed without requesting more tools")
-                                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
-
-                    # Log completion status
-                    if iteration >= max_iterations:
-                        logger.warning(f"Reached maximum iterations ({max_iterations}) for tool execution")
-                        # Add a warning message to the chat
-                        warning_msg = f"Maximum tool execution iterations ({max_iterations}) reached. Stopping tool execution."
-                        new_messages.append(AIMessage(content=warning_msg))
-                    else:
-                        logger.info(f"Tool execution completed after {iteration} iterations")
+                    new_messages, current_completion = self.__perform_tool_calling(completion, messages, llm_client, config)
 
-                    return {"messages": new_messages}
+                    output_msgs = {"messages": new_messages}
+                    if self.output_variables:
+                        if self.output_variables[0] == 'messages':
+                            return output_msgs
+                        output_msgs[self.output_variables[0]] = current_completion.content if current_completion else None
+
+                    return output_msgs
                 else:
                     # Regular text response
                     content = completion.content.strip() if hasattr(completion, 'content') else str(completion)
@@ -273,4 +189,120 @@ class LLMNode(BaseTool):
 
     def _run(self, *args, **kwargs):
         # Legacy support for old interface
-        return self.invoke(kwargs, **kwargs)
+        return self.invoke(kwargs, **kwargs)
+
+    def __perform_tool_calling(self, completion, messages, llm_client, config):
+        # Handle iterative tool-calling and execution
+        new_messages = messages + [completion]
+        iteration = 0
+
+        # Continue executing tools until no more tool calls or max iterations reached
+        current_completion = completion
+        while (hasattr(current_completion, 'tool_calls') and
+               current_completion.tool_calls and
+               iteration < self.steps_limit):
+
+            iteration += 1
+            logger.info(f"Tool execution iteration {iteration}/{self.steps_limit}")
+
+            # Execute each tool call in the current completion
+            tool_calls = current_completion.tool_calls if hasattr(current_completion.tool_calls,
+                                                                  '__iter__') else []
+
+            for tool_call in tool_calls:
+                tool_name = tool_call.get('name', '') if isinstance(tool_call, dict) else getattr(tool_call,
+                                                                                                  'name',
+                                                                                                  '')
+                tool_args = tool_call.get('args', {}) if isinstance(tool_call, dict) else getattr(tool_call,
+                                                                                                  'args',
+                                                                                                  {})
+                tool_call_id = tool_call.get('id', '') if isinstance(tool_call, dict) else getattr(
+                    tool_call, 'id', '')
+
+                # Find the tool in filtered tools
+                filtered_tools = self.get_filtered_tools()
+                tool_to_execute = None
+                for tool in filtered_tools:
+                    if tool.name == tool_name:
+                        tool_to_execute = tool
+                        break
+
+                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)
+
+                        # Create tool message with result - preserve structured content
+                        from langchain_core.messages import ToolMessage
+
+                        # Check if tool_result is structured content (list of dicts)
+                        # TODO: need solid check for being compatible with ToolMessage content format
+                        if isinstance(tool_result, list) and all(
+                                isinstance(item, dict) and 'type' in item for item in tool_result
+                        ):
+                            # Use structured content directly for multimodal support
+                            tool_message = ToolMessage(
+                                content=tool_result,
+                                tool_call_id=tool_call_id
+                            )
+                        else:
+                            # Fallback to string conversion for other tool results
+                            tool_message = ToolMessage(
+                                content=str(tool_result),
+                                tool_call_id=tool_call_id
+                            )
+                        new_messages.append(tool_message)
+
+                    except Exception as e:
+                        logger.error(f"Error executing tool '{tool_name}': {e}")
+                        # Create error tool message
+                        from langchain_core.messages import ToolMessage
+                        tool_message = ToolMessage(
+                            content=f"Error executing {tool_name}: {str(e)}",
+                            tool_call_id=tool_call_id
+                        )
+                        new_messages.append(tool_message)
+                else:
+                    logger.warning(f"Tool '{tool_name}' not found in available tools")
+                    # Create error tool message for missing tool
+                    from langchain_core.messages import ToolMessage
+                    tool_message = ToolMessage(
+                        content=f"Tool '{tool_name}' not available",
+                        tool_call_id=tool_call_id
+                    )
+                    new_messages.append(tool_message)
+
+            # Call LLM again with tool results to get next response
+            try:
+                current_completion = llm_client.invoke(new_messages, config=config)
+                new_messages.append(current_completion)
+
+                # Check if we still have tool calls
+                if hasattr(current_completion, 'tool_calls') and current_completion.tool_calls:
+                    logger.info(f"LLM requested {len(current_completion.tool_calls)} more tool calls")
+                else:
+                    logger.info("LLM completed without requesting more tools")
+                    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
+
+        # Log completion status
+        if iteration >= self.steps_limit:
+            logger.warning(f"Reached maximum iterations ({self.steps_limit}) for tool execution")
+            # Add a warning message to the chat
+            warning_msg = f"Maximum tool execution iterations ({self.steps_limit}) reached. Stopping tool execution."
+            new_messages.append(AIMessage(content=warning_msg))
+        else:
+            logger.info(f"Tool execution completed after {iteration} iterations")
+
+        return new_messages, current_completion
+
+    def __get_struct_output_model(self, llm_client, pydantic_model):
+        return llm_client.with_structured_output(pydantic_model)

alita_sdk/runtime/tools/mcp_inspect_tool.py

@@ -0,0 +1,284 @@
+"""
+MCP Server Inspection Tool.
+Allows inspecting available tools, prompts, and resources on an MCP server.
+"""
+
+import asyncio
+import json
+import logging
+import time
+from typing import Any, Type, Dict, List, Optional
+
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field, ConfigDict
+import aiohttp
+
+logger = logging.getLogger(__name__)
+
+
+class McpInspectInput(BaseModel):
+    """Input schema for MCP server inspection tool."""
+
+    resource_type: str = Field(
+        default="all",
+        description="What to inspect: 'tools', 'prompts', 'resources', or 'all'"
+    )
+
+
+class McpInspectTool(BaseTool):
+    """Tool for inspecting available tools, prompts, and resources on an MCP server."""
+
+    name: str = "mcp_inspect"
+    description: str = "List available tools, prompts, and resources from the MCP server"
+    args_schema: Type[BaseModel] = McpInspectInput
+    return_type: str = "str"
+
+    # MCP server connection details
+    server_name: str = Field(..., description="Name of the MCP server")
+    server_url: str = Field(..., description="URL of the MCP server")
+    server_headers: Optional[Dict[str, str]] = Field(default=None, description="HTTP headers for authentication")
+    timeout: int = Field(default=30, description="Request timeout in seconds")
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def __getstate__(self):
+        """Custom serialization for pickle compatibility."""
+        state = self.__dict__.copy()
+        # Convert headers dict to regular dict to avoid any reference issues
+        if 'server_headers' in state and state['server_headers'] is not None:
+            state['server_headers'] = dict(state['server_headers'])
+        return state
+
+    def __setstate__(self, state):
+        """Custom deserialization for pickle compatibility."""
+        # Initialize Pydantic internal attributes if needed
+        if '__pydantic_fields_set__' not in state:
+            state['__pydantic_fields_set__'] = set(state.keys())
+        if '__pydantic_extra__' not in state:
+            state['__pydantic_extra__'] = None
+        if '__pydantic_private__' not in state:
+            state['__pydantic_private__'] = None
+
+        # Update object state
+        self.__dict__.update(state)
+
+    def _run(self, resource_type: str = "all") -> str:
+        """Inspect the MCP server for available resources."""
+        try:
+            # Always create a new event loop for sync context
+            # This avoids issues with existing event loops in threads
+            import concurrent.futures
+            with concurrent.futures.ThreadPoolExecutor() as executor:
+                future = executor.submit(self._run_in_new_loop, resource_type)
+                return future.result(timeout=self.timeout)
+        except Exception as e:
+            logger.error(f"Error inspecting MCP server '{self.server_name}': {e}")
+            return f"Error inspecting MCP server: {e}"
+
+    def _run_in_new_loop(self, resource_type: str) -> str:
+        """Run the async inspection in a new event loop."""
+        return asyncio.run(self._inspect_server(resource_type))
+
+    async def _inspect_server(self, resource_type: str) -> str:
+        """Perform the actual MCP server inspection."""
+        results = {}
+
+        # Determine what to inspect
+        inspect_tools = resource_type in ["all", "tools"]
+        inspect_prompts = resource_type in ["all", "prompts"]
+        inspect_resources = resource_type in ["all", "resources"]
+
+        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout)) as session:
+
+            # List tools
+            if inspect_tools:
+                try:
+                    tools = await self._list_tools(session)
+                    results["tools"] = tools
+                except Exception as e:
+                    logger.warning(f"Failed to list tools from {self.server_name}: {e}")
+                    results["tools"] = {"error": str(e)}
+
+            # List prompts
+            if inspect_prompts:
+                try:
+                    prompts = await self._list_prompts(session)
+                    results["prompts"] = prompts
+                except Exception as e:
+                    logger.warning(f"Failed to list prompts from {self.server_name}: {e}")
+                    results["prompts"] = {"error": str(e)}
+
+            # List resources
+            if inspect_resources:
+                try:
+                    resources = await self._list_resources(session)
+                    results["resources"] = resources
+                except Exception as e:
+                    logger.warning(f"Failed to list resources from {self.server_name}: {e}")
+                    results["resources"] = {"error": str(e)}
+
+        return self._format_results(results, resource_type)
+
+    def _parse_sse(self, text: str) -> Dict[str, Any]:
+        """Parse Server-Sent Events (SSE) format response."""
+        for line in text.split('\n'):
+            line = line.strip()
+            if line.startswith('data:'):
+                json_str = line[5:].strip()
+                return json.loads(json_str)
+        raise ValueError("No data found in SSE response")
+
+    async def _list_tools(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available tools from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_tools_{int(time.time())}",
+            "method": "tools/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                # Parse SSE format
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    async def _list_prompts(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available prompts from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_prompts_{int(time.time())}",
+            "method": "prompts/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    async def _list_resources(self, session: aiohttp.ClientSession) -> Dict[str, Any]:
+        """List available resources from the MCP server."""
+        request = {
+            "jsonrpc": "2.0",
+            "id": f"list_resources_{int(time.time())}",
+            "method": "resources/list",
+            "params": {}
+        }
+
+        headers = {
+            "Content-Type": "application/json",
+            "Accept": "application/json, text/event-stream",
+            **self.server_headers
+        }
+
+        async with session.post(self.server_url, json=request, headers=headers) as response:
+            if response.status != 200:
+                raise Exception(f"HTTP {response.status}: {await response.text()}")
+
+            # Handle both JSON and SSE responses
+            content_type = response.headers.get('Content-Type', '')
+            if 'text/event-stream' in content_type:
+                text = await response.text()
+                data = self._parse_sse(text)
+            else:
+                data = await response.json()
+
+            if "error" in data:
+                raise Exception(f"MCP Error: {data['error']}")
+
+            return data.get("result", {})
+
+    def _format_results(self, results: Dict[str, Any], resource_type: str) -> str:
+        """Format the inspection results for display."""
+        output_lines = [f"=== MCP Server Inspection: {self.server_name} ==="]
+        output_lines.append(f"Server URL: {self.server_url}")
+        output_lines.append("")
+
+        # Format tools
+        if "tools" in results:
+            if "error" in results["tools"]:
+                output_lines.append(f"❌ TOOLS: Error - {results['tools']['error']}")
+            else:
+                tools = results["tools"].get("tools", [])
+                output_lines.append(f"🔧 TOOLS ({len(tools)} available):")
+                if tools:
+                    for tool in tools:
+                        name = tool.get("name", "Unknown")
+                        desc = tool.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                else:
+                    output_lines.append("  (No tools available)")
+                output_lines.append("")
+
+        # Format prompts
+        if "prompts" in results:
+            if "error" in results["prompts"]:
+                output_lines.append(f"❌ PROMPTS: Error - {results['prompts']['error']}")
+            else:
+                prompts = results["prompts"].get("prompts", [])
+                output_lines.append(f"💬 PROMPTS ({len(prompts)} available):")
+                if prompts:
+                    for prompt in prompts:
+                        name = prompt.get("name", "Unknown")
+                        desc = prompt.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                else:
+                    output_lines.append("  (No prompts available)")
+                output_lines.append("")
+
+        # Format resources
+        if "resources" in results:
+            if "error" in results["resources"]:
+                output_lines.append(f"❌ RESOURCES: Error - {results['resources']['error']}")
+            else:
+                resources = results["resources"].get("resources", [])
+                output_lines.append(f"📁 RESOURCES ({len(resources)} available):")
+                if resources:
+                    for resource in resources:
+                        uri = resource.get("uri", "Unknown")
+                        name = resource.get("name", uri)
+                        desc = resource.get("description", "No description")
+                        output_lines.append(f"  • {name}: {desc}")
+                        output_lines.append(f"    URI: {uri}")
+                else:
+                    output_lines.append("  (No resources available)")
+                output_lines.append("")
+
+        return "\n".join(output_lines)