alita-sdk 0.3.375__py3-none-any.whl → 0.3.417__py3-none-any.whl

alita_sdk/runtime/tools/image_generation.py

@@ -2,16 +2,59 @@
 Image generation tool for Alita SDK.
 """
 import logging
-from typing import Optional, Type, Any
-from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+from typing import Optional, Type, Any, List, Literal
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, Field, create_model, ConfigDict
 
 logger = logging.getLogger(__name__)
 
+name = "image_generation"
+
+
+def get_tools(tools_list: list, alita_client=None, llm=None,
+              memory_store=None):
+    """
+    Get image generation tools for the provided tool configurations.
+
+    Args:
+        tools_list: List of tool configurations
+        alita_client: Alita client instance (required for image generation)
+        llm: LLM client instance (unused for image generation)
+        memory_store: Optional memory store instance (unused)
+
+    Returns:
+        List of image generation tools
+    """
+    all_tools = []
+
+    for tool in tools_list:
+        if (tool.get('type') == 'image_generation' or
+                tool.get('toolkit_name') == 'image_generation'):
+            try:
+                if not alita_client:
+                    logger.error("Alita client is required for image "
+                                 "generation tools")
+                    continue
+
+                toolkit_instance = ImageGenerationToolkit.get_toolkit(
+                    client=alita_client,
+                    toolkit_name=tool.get('toolkit_name', '')
+                )
+                all_tools.extend(toolkit_instance.get_tools())
+            except Exception as e:
+                logger.error(f"Error in image generation toolkit "
+                             f"get_tools: {e}")
+                logger.error(f"Tool config: {tool}")
+                raise
+
+    return all_tools
+
 
 class ImageGenerationInput(BaseModel):
     """Input schema for image generation tool."""
-    prompt: str = Field(description="Text prompt describing the image to generate")
+    prompt: str = Field(
+        description="Text prompt describing the image to generate"
+    )
     n: int = Field(
         default=1, description="Number of images to generate (1-10)",
         ge=1, le=10
@@ -22,7 +65,7 @@ class ImageGenerationInput(BaseModel):
     )
     quality: str = Field(
         default="auto",
-        description="Quality of the generated image ('low', 'medium', 'high', 'auto')"
+        description="Quality of the generated image ('low', 'medium', 'high')"
     )
     style: Optional[str] = Field(
         default=None, description="Style of the generated image (optional)"
@@ -69,7 +112,8 @@ class ImageGenerationTool(BaseTool):
                 else:
                     content_chunks.append({
                         "type": "text",
-                        "text": f"Generated {len(images)} images for prompt: '{prompt}'"
+                        "text": f"Generated {len(images)} images for "
+                                f"prompt: '{prompt}'"
                     })
                 
                 # Add image content for each generated image
@@ -85,7 +129,8 @@ class ImageGenerationTool(BaseTool):
                         content_chunks.append({
                             "type": "image_url",
                             "image_url": {
-                                "url": f"data:image/png;base64,{image_data['b64_json']}"
+                                "url": f"data:image/png;base64,"
+                                       f"{image_data['b64_json']}"
                             }
                         })
                 
@@ -94,7 +139,8 @@ class ImageGenerationTool(BaseTool):
             # Fallback to text response if no images in result
             return [{
                 "type": "text",
-                "text": f"Image generation completed but no images returned: {result}"
+                "text": f"Image generation completed but no images "
+                        f"returned: {result}"
             }]
             
         except Exception as e:
@@ -114,3 +160,53 @@ class ImageGenerationTool(BaseTool):
 def create_image_generation_tool(client):
     """Create an image generation tool with the provided Alita client."""
     return ImageGenerationTool(client=client)
+
+
+class ImageGenerationToolkit(BaseToolkit):
+    """Toolkit for image generation tools."""
+    tools: List[BaseTool] = []
+
+    @staticmethod
+    def toolkit_config_schema() -> BaseModel:
+        """Get the configuration schema for the image generation toolkit."""
+        # Create sample tool to get schema
+        sample_tool = ImageGenerationTool(client=None)
+        selected_tools = {sample_tool.name: sample_tool.args_schema.schema()}
+
+        return create_model(
+            'image_generation',
+            selected_tools=(
+                List[Literal[tuple(selected_tools)]],
+                Field(
+                    default=[],
+                    json_schema_extra={'args_schemas': selected_tools}
+                )
+            ),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Image Generation",
+                    "icon_url": "image_generation.svg",
+                    "hidden": True,
+                    "categories": ["internal_tool"],
+                    "extra_categories": ["image generation"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(cls, client=None, **kwargs):
+        """
+        Get toolkit with image generation tools.
+
+        Args:
+            client: Alita client instance (required)
+            **kwargs: Additional arguments
+        """
+        if not client:
+            raise ValueError("Alita client is required for image generation")
+            
+        tools = [ImageGenerationTool(client=client)]
+        return cls(tools=tools)
+
+    def get_tools(self):
+        return self.tools

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,7 @@ 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')))]
         else:
             # Flow for chat-based LLM node w/o prompt/task from pipeline but with messages in state
             # verify messages structure
@@ -122,14 +123,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 +151,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 +185,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)