alita-sdk 0.3.181__py3-none-any.whl → 0.3.183__py3-none-any.whl

alita_sdk/runtime/tools/llm.py

@@ -1,46 +1,148 @@
+import json
 import logging
 from traceback import format_exc
-from typing import Any, Optional, Dict, List
+from typing import Any, Optional, Dict, List, Union
 
-from langchain_core.messages import HumanMessage, BaseMessage
+from langchain_core.messages import HumanMessage, BaseMessage, SystemMessage, AIMessage
 from langchain_core.tools import BaseTool, ToolException
+from langchain_core.runnables import RunnableConfig
+from pydantic import Field
 
 from ..langchain.utils import _extract_json, create_pydantic_model, create_params
 
 logger = logging.getLogger(__name__)
 
 
-def create_llm_input(prompt: Dict[str, str], params: Dict[str, Any], kwargs: Dict[str, Any]) -> list[BaseMessage]:
-    logger.info(f"Creating LLM input with prompt: {prompt}, params: {params}, kwargs: {kwargs}")
-    if prompt.get('type') == 'fstring' and params:
+def create_llm_input_with_messages(
+    prompt: Dict[str, str], 
+    messages: List[BaseMessage], 
+    params: Dict[str, Any]
+) -> List[BaseMessage]:
+    """
+    Create LLM input by combining system prompt with chat history messages.
+    
+    Args:
+        prompt: The prompt configuration with template
+        messages: List of chat history messages
+        params: Additional parameters for prompt formatting
+        
+    Returns:
+        List of messages to send to LLM
+    """
+    logger.info(f"Creating LLM input with messages: {len(messages)} messages, params: {params}")
+    
+    # Build the input messages
+    input_messages = []
+    
+    # Add system message from prompt if available
+    if prompt and prompt.get('template'):
         try:
-            return [HumanMessage(content=prompt['value'].format(**params))]
+            system_content = prompt['template'].format(**params) if params else prompt['template']
+            input_messages.append(SystemMessage(content=system_content))
         except KeyError as e:
-            error_msg = f"KeyError in input formatting - make sure you have added all required state variables as input to the node: {e}"
+            error_msg = f"KeyError in prompt formatting: {e}. Available params: {list(params.keys())}"
             logger.error(error_msg)
             raise ToolException(error_msg)
-    else:
-        return kwargs.get("messages") + [HumanMessage(prompt['value'])]
+    
+    # Add the chat history messages
+    if messages:
+        input_messages.extend(messages)
+    
+    return input_messages
 
 
 class LLMNode(BaseTool):
-    name: str = 'LLMNode'
-    prompt: Dict[str, str]
-    description: str = 'This is tool node for LLM'
-    client: Any = None
-    return_type: str = "str"
-    response_key: str = "messages"
-    structured_output_dict: Optional[dict[str, str]] = None
-    output_variables: Optional[List[str]] = None
-    input_variables: Optional[List[str]] = None
-    structured_output: Optional[bool] = False
+    """Enhanced LLM node with chat history and tool binding support"""
+    
+    # Override BaseTool required fields
+    name: str = Field(default='LLMNode', description='Name of the LLM node')
+    description: str = Field(default='This is tool node for LLM with chat history and tool support', description='Description of the LLM node')
+    
+    # LLM-specific fields
+    prompt: Dict[str, str] = Field(default_factory=dict, description='Prompt configuration')
+    client: Any = Field(default=None, description='LLM client instance')
+    return_type: str = Field(default="str", description='Return type')
+    response_key: str = Field(default="messages", description='Response key')
+    structured_output_dict: Optional[dict[str, str]] = Field(default=None, description='Structured output dictionary')
+    output_variables: Optional[List[str]] = Field(default=None, description='Output variables')
+    input_variables: Optional[List[str]] = Field(default=None, description='Input variables')
+    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')
 
-    def _run(self, *args, **kwargs):
-        params = create_params(self.input_variables, kwargs)
-        logger.info(f"LLM Node params: {params}")
-        llm_input = create_llm_input(self.prompt, params, kwargs)
+    def get_filtered_tools(self) -> List[BaseTool]:
+        """
+        Filter available tools based on tool_names list.
+        
+        Returns:
+            List of filtered tools
+        """
+        if not self.available_tools:
+            return []
+            
+        if not self.tool_names:
+            # If no specific tool names provided, return all available tools
+            return self.available_tools
+            
+        # Filter tools by name
+        filtered_tools = []
+        available_tool_names = {tool.name: tool for tool in self.available_tools}
+        
+        for tool_name in self.tool_names:
+            if tool_name in available_tool_names:
+                filtered_tools.append(available_tool_names[tool_name])
+                logger.debug(f"Added tool '{tool_name}' to LLM node")
+            else:
+                logger.warning(f"Tool '{tool_name}' not found in available tools: {list(available_tool_names.keys())}")
+        
+        return filtered_tools
+
+    def invoke(
+        self,
+        state: Union[str, dict],
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> dict:
+        """
+        Invoke the LLM node with proper message handling and tool binding.
+        
+        Args:
+            state: The current state containing messages and other variables
+            config: Optional runnable config
+            **kwargs: Additional keyword arguments
+            
+        Returns:
+            Updated state with LLM response
+        """
+        # Extract messages from state
+        
+        messages = state.get("messages", []) if isinstance(state, dict) else []
+        logger.info(f"Invoking LLMNode with {len(messages)} messages")
+        logger.info("Messages: %s", messages)
+        # Create parameters for prompt formatting from state
+        params = {}
+        if isinstance(state, dict):
+            for var in self.input_variables or []:
+                if var != "messages" and var in state:
+                    params[var] = state[var]
+        
+        # Create LLM input with proper message handling
+        llm_input = create_llm_input_with_messages(self.prompt, messages, params)
+        
+        # Get the LLM client, potentially with tools bound
+        llm_client = self.client
+        
+        if len(self.tool_names or []) > 0:
+            filtered_tools = self.get_filtered_tools()
+            if filtered_tools:
+                logger.info(f"Binding {len(filtered_tools)} tools to LLM: {[t.name for t in filtered_tools]}")
+                llm_client = self.client.bind_tools(filtered_tools)
+            else:
+                logger.warning("No tools to bind to LLM")
+        
         try:
-            if self.structured_output and len(self.output_variables) > 0:
+            if self.structured_output and self.output_variables:
+                # Handle structured output
                 struct_params = {
                     key: {
                         "type": 'list[str]' if 'list' in value else value,
@@ -48,27 +150,144 @@ class LLMNode(BaseTool):
                     }
                     for key, value in (self.structured_output_dict or {}).items()
                 }
-                stuct_model = create_pydantic_model(f"LLMOutput", struct_params)
-                llm = self.client.with_structured_output(stuct_model)
-                completion = llm.invoke(llm_input)
+                struct_model = create_pydantic_model(f"LLMOutput", struct_params)
+                llm = llm_client.with_structured_output(struct_model)
+                completion = llm.invoke(llm_input, 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'])}]
+                
                 return result
             else:
-                completion = self.client.invoke(llm_input)
-                result = completion.content.strip()
-                response = _extract_json(result) or {}
-                response_data = {key: response[key] for key in response if key in self.output_variables}
-                if not response_data.get('messages'):
-                    response_data['messages'] = [
-                        {"role": "assistant", "content": response_data.get(self.response_key) or result}]
-                return response_data
-        except ValueError:
-            if self.output_variables:
-                return {self.output_variables[0]: result, "messages": [{"role": "assistant", "content": result}]}
-            else:
-                return {"messages": [{"role": "assistant", "content": result}]}
+                # Handle regular completion
+                completion = llm_client.invoke(llm_input, config=config)
+                logger.info(f"Initial completion: {completion}")
+                # 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
+                                    from langchain_core.messages import ToolMessage
+                                    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")
+                    
+                    return {"messages": new_messages}
+                else:
+                    # Regular text response
+                    content = completion.content.strip() if hasattr(completion, 'content') else str(completion)
+                    
+                    # Try to extract JSON if output variables are specified (but exclude 'messages' which is handled separately)
+                    json_output_vars = [var for var in (self.output_variables or []) if var != 'messages']
+                    if json_output_vars:
+                        try:
+                            response = _extract_json(content) or {}
+                            response_data = {key: response.get(key) for key in json_output_vars if key in response}
+                            
+                            # Always add the messages to the response
+                            new_messages = messages + [AIMessage(content=content)]
+                            response_data['messages'] = new_messages
+                            
+                            return response_data
+                        except (ValueError, json.JSONDecodeError) as e:
+                            # LLM returned non-JSON content, treat as plain text
+                            logger.warning(f"Expected JSON output but got plain text. Output variables specified: {json_output_vars}. Error: {e}")
+                            # Fall through to plain text handling
+                    
+                    # Simple text response (either no output variables or JSON parsing failed)
+                    new_messages = messages + [AIMessage(content=content)]
+                    return {"messages": new_messages}
+                        
         except Exception as e:
             logger.error(f"Error in LLM Node: {format_exc()}")
-            return {"messages": [{"role": "assistant", "content": f"Error: {e}"}]}
+            error_msg = f"Error: {e}"
+            new_messages = messages + [AIMessage(content=error_msg)]
+            return {"messages": new_messages}
+
+    def _run(self, *args, **kwargs):
+        # Legacy support for old interface
+        return self.invoke(kwargs, **kwargs)

alita_sdk/runtime/tools/mcp_server_tool.py

@@ -35,6 +35,21 @@ class McpServerTool(BaseTool):
                 field_type = float
             elif field_type == 'boolean':
                 field_type = bool
+            elif field_type == 'object':#Dict[str, Any]
+                nested_model = McpServerTool.create_pydantic_model_from_schema(field_info)
+                field_type = nested_model
+            elif field_type == 'array':
+                item_schema = field_info['items']
+                item_type = McpServerTool.create_pydantic_model_from_schema(item_schema) if item_schema['type'] == 'object' else (
+                    str if item_schema['type'] == 'string' else
+                    int if item_schema['type'] == 'integer' else
+                    float if item_schema['type'] == 'number' else
+                    bool if item_schema['type'] == 'boolean' else
+                    None
+                )
+                if item_type is None:
+                    raise ValueError(f"Unsupported array item type: {item_schema['type']}")
+                field_type = list[item_type]
             else:
                 raise ValueError(f"Unsupported field type: {field_type}")
 

alita_sdk/runtime/utils/streamlit.py

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)
 
 ai_icon = b'<plain_txt_msg:img>iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAAAXNSR0IArs4c6QAAAMZlWElmTU0AKgAAAAgABgESAAMAAAABAAEAAAEaAAUAAAABAAAAVgEbAAUAAAABAAAAXgEoAAMAAAABAAIAAAExAAIAAAAVAAAAZodpAAQAAAABAAAAfAAAAAAAAABIAAAAAQAAAEgAAAABUGl4ZWxtYXRvciBQcm8gMy41LjEAAAAEkAQAAgAAABQAAACyoAEAAwAAAAEAAQAAoAIABAAAAAEAAAAgoAMABAAAAAEAAAAgAAAAADIwMjQ6MDQ6MDMgMTk6NDA6NDQATjJeeQAAAAlwSFlzAAALEwAACxMBAJqcGAAAA7BpVFh0WE1MOmNvbS5hZG9iZS54bXAAAAAAADx4OnhtcG1ldGEgeG1sbnM6eD0iYWRvYmU6bnM6bWV0YS8iIHg6eG1wdGs9IlhNUCBDb3JlIDYuMC4wIj4KICAgPHJkZjpSREYeG1sbnM6cmRmPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5LzAyLzIyLXJkZi1zeW50YXgtbnMjIj4KICAgICAgPHJkZjpEZXNjcmlwdGlvbiByZGY6YWJvdPC9IgogICAgICAgICAgICAgeG1sbnM6dGlmZj0iaHR0cDovL25zLmFkb2JlLmNvbS90aWZmLzEuMC8iCiAgICAgICAgICAgIHhtbG5zOmV4aWY9X2h0dHA6Ly9ucy5hZG9iZS5jb20vZXhpZi8xLjAvIgogICAgICAgICAgICB4bWxuczp4bXA9Imh0dHA6Ly9ucy5hZG9iZS5jb20veGFwLzEuMC8iPgogICAgICAgICA8dGlmZjpZUmVzb2x1dGlvbj43MjAwMDAvMTAwMDA8L3RpZmY6WVJlc29sdXRpb24+CiAgICAgICAgIDx0aWZmOlhSZXNvbHV0aW9uPjcyMDAwMC8xMDAwMDwvdGlmZjpYUmVzb2x1dGlvbj4KICAgICAgICAgPHRpZmY6UmVzb2x1dGlvblVuaXQ+MjwvdGlmZjpSZXNvbHV0aW9uVW5pdD4KICAgICAgICAgPHRpZmY6T3JpZW50YXRpb24+MTwvdGlmZjpPcmllbnRhdGlvbj4KICAgICAgICAgPGV4aWY6UGl4ZWxZRGltZW5zaW9uPjMyPC9leGlmOlBpeGVsWURpbWVuc2lvbj4KICAgICAgICAgPGV4aWY6UGl4ZWxYRGltZW5zaW9uPjMyPC9leGlmOlBpeGVsWERpbWVuc2lvbj4KICAgICAgICAgPHhtcDpNZXRhZGF0YURhdGU+MjAyNC0wNC0wM1QxOTo0Mjo1OSswMzowMDwveG1wOk1ldGFkYXRhRGF0ZT4KICAgICAgICAgPHhtcDpDcmVhdGVEYXRlPjIwMjQtMDQtMDNUMTk6NDI6MjMrMDM6MDA8L3htcDpDcmVhdGVEYXRlPgogICAgICAgICA8eG1wOkNyZWF0b3JUb29sPlBpeGVsbWF0b3IgUHJvIDMuNS4xPC94bXA6Q3JlYXRvclRvb2w+CiAgICAgIDwvcmRmOkRlc2NyaXB0aW9uPgogICA8L3JkZjpSREY+CjwveDp4bXBtZXRhPgrdpg3+AAAH30lEQVRYCe1Xe4yU1RU/936P+eab187s7izIo4LIo1IR12AJpRIphNIlrVbQtqSVNgVaGw0+Uki02SiYJlaSltoE0kSJFRUo6m4LBWxdwdZQqSWuPOxSyurCzj5nZ+ab+d739txvd8YBFtz/25vcvWfuveec3/mdc+/9FuB/vZExEkD+ef9rCcfkadM358xOTmvMOYWpJddMUiL5UTWccbh35kKp91hSSXa4EX/wth0rTLTNP8v+ZwEgh7+7e9IErX6pRpVl2BeqkpoeVhJ/L7PPCc+5xbMc+KG8Zxx485P2tscOPVa8FohrAjj6w5amcXr6kRBR50mU6tcyVL3GEYHP/UzRM4+cGux44p5dazqq16vlqwEgr31v18Y5DTdt4QSutqfazlVlBNJ3sueD1XftvP/QaJuuMN66trUuGYs3p8LJBy5RYByo6YNkuCAPOSAVPSAuAwGPhSTw4wp4CRWYLgPTpIoqYR7I/R1d+WzXo417Hnm1sjAifLoTJ7avbdXrk/FHY5HEel8iqo+rOAJzfaC9JsgZE+iABdz2wUeefXQuOvOR85IHkHcATC+Yc3QEpVBgxIOY0R0PcZi37LrG0y+dffvf1SBo9Y9EnC4Jx2seclUpaisELOyu6QJ0FcFH4xbuLkVlKMVG7yY6tZEpJkBmSmATBj51QOYGaAqdPKt24s59izdNr/ZZYWDl1t2p20OTWmVVrRdR+5QAR8pJnwU+argqBTf0aXdC4DkKsRwVPE+TJDckkco6Ru4jEO74EJUyELGzgU88stGYFK+7Ib94/8H8QaQMQA5WAOjEL8x9PH+RTpSGrGCKIM1y0QEvjGFrtFyJ3PPc/3iu9XfMQEfJMXoY4XKNmpzsEzZbVcO3U0oT4nCK4iLIQLx0AfCqAM7RBsXU2TWLZo5LLoIuOCAcBQDWth6fHkqk7roALmglG6iHBYc5dUUxabgrOO7csUxjr2saz/2t6+CH2w5sywsDojVDM629b9K4VKLui+F43RZV1WYKACm/BzTbQOcSEATAGQWrmGwIy+7y1sbtb6/4x7pSAICE9Tu5QtOFeAiyWMnRPhMIRl3dsvn+3+7auvmhNmgLqKteQwAMXoGLOLdv2rRpf3xi9Qt/1rXoglmFcxi9hwCQDQRQyNWBw0KUcWcJeMmf4/4Snb9ha1ivr78FVDks8txfp0MxLFUKrRiVWR8Yv1v79OjOq4EI+ezZs3avM/DjBOtoj5JsQL9IgcsU6B2aABxrS1cj03Ikf73YT2es+moDV+QpXCKESxSMRBiG6vUKgILmn7roZp+EUSIXBkZrc/Jv9M8oHe9WKNYT5p0Bhd78BPCIDOgHQVApGapZKHRlSZaTmJp6hsiIxMFDFoYaYhCVsQ48n1klc2/vJxfPj+boanNTY/KSeMj5EvrGewRZNcdD1kpjEeKdggAIdsTweaFPfdsNeczTh5EhOnEEVQmKtXEwNDKQhdKJHTvW4WUwttbykxVTEjH6lCxzHdMOA9Z4uFC4EY91OXrBAN4x3KkVFikjhICMB0ZQIzouipEhchaJSjRWO2VF8/YxPUQdT35z7oLra/fpGplk4/HpLM6Cjtyt4RIFBMNBD2yjD5kEVU6JpjigKLZwXNkgZNwIspLSI8lNc+Z+Zcf6X+2fi4DF6bqiNWMgJ5/+2sJ0KvScHia3ZJzPQXv+y3C+NBvzrlSCCgIUQWKXJCU4xlRPpg2ihbLVi4EsmMBOVDmtRePfmXzjrceaXz7d8oPHX1wG0KiUUSxCmPf+YukD4+rqWrr5nPmHBr8P7+WboN+dgG+ChM7wGgm6CBDlkUAjWuy8sCGffuP3PdO/8fWPqRYJioPg6yKKRRRK0FlZpooWTzbNbFzc9Owrd7hWNnuO2s7QdWpp6gnQ6i0jAcQGkPDQS1h9w/p4++FDFchYkESkleELinfSOaPzrwGAtg1rcjcsz3zIwhGbUhKqMIGOAxkBVeZGZDw7SiSRmiHhc1x0GFg4ShS7xPC2wzdA0Cz20ioZT1jABAKxidOtRWLtAoAoBG4N9h/hhOTKhShGUQ/CUDAnaBuRq+cr8mXrlVoq64zYK9uwqf+OyWCwDABemj/7XdvIvRc4FI7LzssjGqjMEXzjCM+4zH0/m8u0ZHO9B/BB+ogBG8QcswpbQqdKb3iegkdZwSL+H9a9uSoowvJryIqdnU+pM/UFVFJqgs1l6nEUyCnWk2e753zLeLl4oXPf+/u2tLe1Db8LK1c2R+fefMeiWCh5rwahe6hEtOEUjIAo28B6wK/nY31m/1uCecFA9bEia052bYzUpTdLLqcS5lbkOBhRZnnjjNnf+yO/u/vdbQ8ux3K7sjVv2J3So+ObUlrNNpXL8bJu2Q51fD8z+HHT+j13/6msXQ0Avn30g6TekN6pR1MrJPyYQCABAG4UB/71l/23vb5pzfmy4rXGZ3761uraWP321ad6dSCd3Wc2PfzqfeIVrLTgNir/2rXw5qw3lP+ZY5WODBcM3uXg9eWGer81VufCVrHQ97rjmjt8yl1RkFgzRtbJ/vJy52LvJQDExAvzpp9w+noeLGUH9/qc2aVS4fnBXOEdsTbW1vybVUbfYNevHd/9CL+U8jkr+8yQZ2weTf+SFFRvWHn4eCJC1aWcSKd23nnTyeq1Mcrk2Y2H73aKViGUMY8+vGeV+Fft/+0KBv4LJG7QdeOMt6wAAAAASUVORK5CYII=<!plain_txt_msg>'
 user_icon = b'<plain_txt_msg:img>iVBORw0KGgoAAAANSUhEUgAAACAAAAAgCAYAAABzenr0AAAAAXNSR0IArs4c6QAAAMZlWElmTU0AKgAAAAgABgESAAMAAAABAAEAAAEaAAUAAAABAAAAVgEbAAUAAAABAAAAXgEoAAMAAAABAAIAAAExAAIAAAAVAAAAZodpAAQAAAABAAAAfAAAAAAAAABIAAAAAQAAAEgAAAABUGl4ZWxtYXRvciBQcm8gMy41LjEAAAAEkAQAAgAAABQAAACyoAEAAwAAAAEAAQAAoAIABAAAAAEAAAAgoAMABAAAAAEAAAAgAAAAADIwMjQ6MDQ6MDMgMTk6NDI6MjMAfz7nbAAAAAlwSFlzAAALEwAACxMBAJqcGAAAA7BpVFh0WE1MOmNvbS5hZG9iZS54bXAAAAAAADx4OnhtcG1ldGEgeG1zbnM6eD0iYWRvYmU6bnM6bWV0YS8iIHg6eG1wdGs9IlhNUCBDb3JlIDYuMC4wIj4KICAgPHJkZjpSREYgeG1zbnM6cmRmPSJodHRwOi8vd3d3LnczLm9yZy8xOTk5LzAyLzIyLXJkZi1zeW50YXgtbnMjIj4KICAgICAgPHJkZjpEZXNjcmlwdGlvbiByZGY6YWJvdXQ9IiIKICAgICAgICAgICAgeG1zbnM6dGlmZj0iaHR0cDovL25zLmFkb2JlLmNvbS90aWZmLzEuMC8iCiAgICAgICAgICAgIHhtbG5zOmV4aWY9Imh0dHA6Ly9ucy5hZG9iZS5jb20vZXhpZi8xLjAvIgogICAgICAgICAgICB4bWxuczp4bXA9Imh0dHA6Ly9ucy5hZG9iZS5jb20veGFwLzEuMC8iPgogICAgICAgICA8dGlmZjpZUmVzb2x1dGlvbj43MjAwMDAvMTAwMDA8L3RpZmY6WVJlc29sdXRpb24+CiAgICAgICAgIDx0aWZmOlhSZXNvbHV0aW9uPjcyMDAwMC8xMDAwMDwvdGlmZjpYUmVzb2x1dGlvbj4KICAgICAgICAgPHRpZmY6UmVzb2x1dGlvblVuaXQ+MjwvdGlmZjpSZXNvbHV0aW9uVW5pdD4KICAgICAgICAgPHRpZmY6T3JpZW50YXRpb24+MTwvdGlmZjpPcmllbnRhdGlvbj4KICAgICAgICAgPGV4aWY6UGl4ZWxZRGltZW5zaW9uPjMyPC9leGlmOlBpeGVsWURpbWVuc2lvbj4KICAgICAgICAgPGV4aWY6UGl4ZWxYRGltZW5zaW9uPjMyPC9leGlmOlBpeGVsWERpbWVuc2lvbj4KICAgICAgICAgPHhtcDpNZXRhZGF0YURhdGU+MjAyNC0wNC0wM1QxOTo0Mjo1OSswMzowMDwveG1wOk1ldGFkYXRhRGF0ZT4KICAgICAgICAgPHhtcDpDcmVhdGVEYXRlPjIwMjQtMDQtMDNUMTk6NDI6MjMrMDM6MDA8L3htcDpDcmVhdGVEYXRlPgogICAgICAgICA8eG1wOkNyZWF0b3JUb29sPlBpeGVsbWF0b3IgUHJvIDMuNS4xPC94bXA6Q3JlYXRvclRvb2w+CiAgICAgIDwvcmRmOkRlc2NyaXB0aW9uPgogICA8L3JkZjpSREY+CjwveDp4bXBtZXRhPgorTryeAAAFAUlEQVRYCc1XbWgcRRie2Zndy6XGpD+ijbZWxT8G7QfKnWlremk0JpriR7jUVsUgiGK1EhQKWnGFKiqUSLTJBbRJG6r2IuJHkNgkJrZJSxIQLWiJloI/mhgLirnE3O19jM/s7SZXLnd7VwRduHtnZ96PZ96d9312CfmPL5pvfN0f1NisaI6ovCHi4uVhzgqjLjaD+7GYqrV2B7Z+k4/PvAAc2Ny4KqapxyIqKzU0TiIaww9SXRobKjuhMF7b07JpIRcgSi5KUqfdc2SbQuhxDEvlvSAkIihpJVTsxW2fnJMX5ioXGDmv6yIn3zxplv2/bcvRlQnD6KeEmU4Zide83ts4QBDdsnzbpw/xVTP0C9zX4bfqlDE+ArnJWs8ockJJotGX4CEZPCZqXz6+qz8luOl8WK+Kfdzuu5dQ+i2ygAzRCoAqyBjZWnAEMOQb4tjnC1IffodePPnI19mcMpf7fqzHZGoSpLglm65ccwTwc+jiCiGfbPK4Bp0c9rx1+1/Y/XRSTzzgpO8IIMoWFs8JjZMLTg7NdUHsCli0zWTnCECNu2O2sWDkWnucVVLitjK2aJtJ3xHAdFHpPPJvnnb87cjkyJ737+0vpkSUJe/FZ/Z8JukIQB+uiglBDkgHOAm+N6uO3pPJmZyPRQo+hzBTr5BQczZdueYIwHTgUt9AXSXMOqC077Wa4N0mnBTvsg88/MzwV8jVVnNakFGUZjhFZdlhzq34vYquaoOxPkNjPNl6eQTtuCOs8QuGplRhrtZszbItu9i0p7Bita5TgM5+5QxAumnxdfsiCg9KLrB5wOQE1eKEJDcMF5esruvSb3DcvfSZFwBpINmQSjbUeIOhqeURTXFj978jM2MRrrR+FKjOiw1zOwMyMq6DvuAVBX/En0S7vRPYr0KDUoWgskm5USFrBWU7Hnx+aEtSO7f/nDLQ6essmA+TQJRzP1JfaFKwScPJ1GP3l1Kyxs+FNaV5cL+31wmGI4CA54P1CaKAYGgxnj0xgyE45ACAjOKFZM5w8etxGLfjXeC6VHDQOX3y1duyMmJWAO2eQ/Xo61/KXchMI8Ck4ebbFvonf9OJnnbCZSmWzGqPGkx5H5liSTBsaiKybg3JUBEZAbR5O+9A7Z82gxMRw3PeuXvsiU+cUmqvV+8b/x6Hc72slqjGzpaELq6TlG2v23LZQ9jqbb0SwQcspagQyoZ8gku7wf2eDUShPbLQEoTePLOy7F07aKpMAxD0BxkXRR9CaYVUROrrn51o+jHVKNdx2ZlfdgoqvjP9KPTptR0/WRyx5CENQGgqVIjlSkslvHusSb4HXtbV09MYxxlqksaSzWhCe0yOU680AEXXFP0NhROWUsFBb1dNqkE+Yz+ySYU4bNsIxei2x7ZMA9AI1DEa2gWFeakEB70B76Fy2yAfOXXrTUF0qY3SRkmIwK9PlVtvSkte0gDIpT1je2Zxdu6y1FTk8Yd2b1fDkpnzqHrfxBmU2EOW5tmr/5x+bjmrjGUolds8Xffh6ZndLKGgD3A2GXXoA8UhtQms2YEeoMgSRO+4vD5go7U7IQBYndD6ElLZIGh3BIQ0F1bZjYaL1aPu15jNx2JHdMhTI69s3Gz7Wk5mzYBtILlgTnKByv1ouYWLn2WXfJItfZ5h5+cM9V/iAhuElDrYUGXxx6MuXgdeuAU7LsNPRSZmAeo8xuNRlXZ/+k7laKrd/3r8D5cLzopAT7EBAAAAAElFTkSuQmCC<!plain_txt_msg>'
-agent_types = ["pipeline", "react", "xml", "openai"]
+agent_types = ["pipeline", "react", "xml", "openai", "predict"]
 
 def img_to_txt(filename):
     msg = b"<plain_txt_msg:img>"
@@ -1034,18 +1034,14 @@ def run_streamlit(st, ai_icon=None, user_icon=None):
             st.session_state.messages.append({"role": "user", "content": prompt})
             with st.chat_message("assistant", avatar=ai_icon):
                 st_cb = AlitaStreamlitCallback(st)
+                logger.info(st.session_state.messages)
                 response = st.session_state.agent_executor.invoke(
                     {"input": prompt, "chat_history": st.session_state.messages[:-1]},
                     { 'callbacks': [st_cb], 'configurable': {"thread_id": st.session_state.thread_id}}
                 )
-                try:
-                    output =response["output"]
-                except KeyError:
-                    logger.info(response)
-                    output = response['messages'][0].content
-                st.write(output)
+                st.write(response["output"])
                 st.session_state.thread_id = response.get("thread_id", None)
-                st.session_state.messages.append({"role": "assistant", "content": output})
+                st.session_state.messages.append({"role": "assistant", "content": response["output"]})
     
     elif st.session_state.llm and st.session_state.show_toolkit_testing and st.session_state.configured_toolkit:
         # Toolkit Testing Main View

alita_sdk/tools/__init__.py

@@ -79,6 +79,7 @@ _safe_import_tool('pptx', 'pptx', 'get_tools', 'PPTXToolkit')
 _safe_import_tool('postman', 'postman', 'get_tools', 'PostmanToolkit')
 _safe_import_tool('memory', 'memory', 'get_tools', 'MemoryToolkit')
 _safe_import_tool('zephyr_squad', 'zephyr_squad', 'get_tools', 'ZephyrSquadToolkit')
+_safe_import_tool('slack', 'slack', 'get_tools', 'SlackToolkit')
 
 # Log import summary
 available_count = len(AVAILABLE_TOOLS)

alita_sdk/tools/ado/test_plan/test_plan_wrapper.py

@@ -10,7 +10,9 @@ from langchain_core.tools import ToolException
 from msrest.authentication import BasicAuthentication
 from pydantic import create_model, PrivateAttr, model_validator, SecretStr
 from pydantic.fields import FieldInfo as Field
+import xml.etree.ElementTree as ET
 
+from ..work_item import AzureDevOpsApiWrapper
 from ...elitea_base import BaseToolApiWrapper
 
 logger = logging.getLogger(__name__)
@@ -101,6 +103,16 @@ TestCaseAddModel = create_model(
     suite_id=(int, Field(description="ID of the test suite to which test cases are to be added"))
 )
 
+TestCaseCreateModel = create_model(
+    "TestCaseCreateModel",
+    project=(str, Field(description="Project ID or project name")),
+    plan_id=(int, Field(description="ID of the test plan to which test cases are to be added")),
+    suite_id=(int, Field(description="ID of the test suite to which test cases are to be added")),
+    title=(str, Field(description="Test case title")),
+    description=(str, Field(description="Test case description")),
+    test_steps=(str, Field(description="""Json array with test steps. Example: [{"action": "Some action", "expectedResult": "Some expectation"},...]""")),
+)
+
 TestCaseGetModel = create_model(
     "TestCaseGetModel",
     project=(str, Field(description="Project ID or project name")),
@@ -202,17 +214,51 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
             logger.error(f"Error getting test suite(s): {e}")
             return ToolException(f"Error getting test suite(s): {e}")
 
-    def add_test_case(self, suite_test_case_create_update_parameters: str, project: str, plan_id: int, suite_id: int):
+    def add_test_case(self, suite_test_case_create_update_parameters, project: str, plan_id: int, suite_id: int):
         """Add a test case to a suite in Azure DevOps."""
         try:
-            params = json.loads(suite_test_case_create_update_parameters)
-            suite_test_case_create_update_params_obj = [SuiteTestCaseCreateUpdateParameters(**param) for param in params]
+            if isinstance(suite_test_case_create_update_parameters, str):
+                suite_test_case_create_update_parameters = json.loads(suite_test_case_create_update_parameters)
+            suite_test_case_create_update_params_obj = [SuiteTestCaseCreateUpdateParameters(**param) for param in suite_test_case_create_update_parameters]
             test_cases = self._client.add_test_cases_to_suite(suite_test_case_create_update_params_obj, project, plan_id, suite_id)
             return [test_case.as_dict() for test_case in test_cases]
         except Exception as e:
             logger.error(f"Error adding test case: {e}")
             return ToolException(f"Error adding test case: {e}")
 
+    def create_test_case(self, project: str, plan_id: int, suite_id: int, title: str, description: str, test_steps: str):
+        """Creates a new test case in specified suite in Azure DevOps."""
+        work_item_wrapper = AzureDevOpsApiWrapper(organization_url=self.organization_url, token=self.token.get_secret_value(), project=project)
+        work_item_json = self.build_ado_test_case(title, description, json.loads(test_steps))
+        created_work_item_id = work_item_wrapper.create_work_item(work_item_json=json.dumps(work_item_json), wi_type="Test Case")['id']
+        return self.add_test_case([{"work_item":{"id":created_work_item_id}}], project, plan_id, suite_id)
+
+    def build_ado_test_case(self, title, description, steps):
+        """
+        :param title: test title
+        :param description: test description
+        :param steps: steps [(action, expected result), ...]
+        :return: JSON with ADO fields
+        """
+        steps_elem = ET.Element("steps")
+
+        for idx, step in enumerate(steps, start=1):
+            step_elem = ET.SubElement(steps_elem, "step", id=str(idx), type="Action")
+            action_elem = ET.SubElement(step_elem, "parameterizedString", isformatted="true")
+            action_elem.text = step["action"]
+            expected_elem = ET.SubElement(step_elem, "parameterizedString", isformatted="true")
+            expected_elem.text = step["expectedResult"]
+
+        steps_xml = ET.tostring(steps_elem, encoding="unicode")
+
+        return {
+            "fields": {
+                "System.Title": title,
+                "System.Description": description,
+                "Microsoft.VSTS.TCM.Steps": steps_xml
+            }
+        }
+
     def get_test_case(self, project: str, plan_id: int, suite_id: int, test_case_id: str):
         """Get a test case from a suite in Azure DevOps."""
         try:
@@ -280,6 +326,12 @@ class TestPlanApiWrapper(BaseToolApiWrapper):
                 "args_schema": TestCaseAddModel,
                 "ref": self.add_test_case,
             },
+            {
+                "name": "create_test_case",
+                "description": self.create_test_case.__doc__,
+                "args_schema": TestCaseCreateModel,
+                "ref": self.create_test_case,
+            },
             {
                 "name": "get_test_case",
                 "description": self.get_test_case.__doc__,

alita_sdk/tools/ado/work_item/ado_wrapper.py

@@ -150,14 +150,15 @@ class AzureDevOpsApiWrapper(BaseToolApiWrapper):
 
         return parsed_items
 
-    def _transform_work_item(self, work_item_json: str):
+    def _transform_work_item(self, work_item_json):
         try:
             # Convert the input JSON to a Python dictionary
-            params = json.loads(work_item_json)
+            if isinstance(work_item_json, str):
+                work_item_json = json.loads(work_item_json)
         except (json.JSONDecodeError, ValueError) as e:
             raise ToolException(f"Issues during attempt to parse work_item_json: {e}")
 
-        if 'fields' not in params:
+        if 'fields' not in work_item_json:
             raise ToolException("The 'fields' property is missing from the work_item_json.")
 
             # Transform the dictionary into a list of JsonPatchOperation objects
@@ -167,11 +168,11 @@ class AzureDevOpsApiWrapper(BaseToolApiWrapper):
                 "path": f"/fields/{field}",
                 "value": value
             }
-            for field, value in params["fields"].items()
+            for field, value in work_item_json["fields"].items()
         ]
         return patch_document
 
-    def create_work_item(self, work_item_json: str, wi_type="Task"):
+    def create_work_item(self, work_item_json, wi_type="Task"):
         """Create a work item in Azure DevOps."""
         try:
             patch_document = self._transform_work_item(work_item_json)
@@ -185,7 +186,10 @@ class AzureDevOpsApiWrapper(BaseToolApiWrapper):
                 project=self.project,
                 type=wi_type
             )
-            return f"Work item {work_item.id} created successfully. View it at {work_item.url}."
+            return {
+                "id": work_item.id,
+                "message": f"Work item {work_item.id} created successfully. View it at {work_item.url}."
+            }
         except Exception as e:
             if "unknown value" in str(e):
                 logger.error(f"Unable to create work item due to incorrect assignee: {e}")

alita_sdk/tools/carrier/api_wrapper.py

@@ -70,6 +70,9 @@ class CarrierAPIWrapper(BaseModel):
     def get_integrations(self, name: str):
         return self._client.get_integrations(name)
 
+    def get_available_locations(self):
+        return self._client.get_available_locations()
+
     def run_test(self, test_id: str, json_body):
         return self._client.run_test(test_id, json_body)