alita-sdk 0.3.257__py3-none-any.whl → 0.3.584__py3-none-any.whl

alita_sdk/runtime/tools/function.py

@@ -1,18 +1,33 @@
+import json
 import logging
+from copy import deepcopy
 from json import dumps
 
 from langchain_core.callbacks import dispatch_custom_event
 from langchain_core.messages import ToolCall
 from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import BaseTool
-from typing import Any, Optional, Union, Annotated
+from langchain_core.tools import BaseTool, ToolException
+from typing import Any, Optional, Union
 from langchain_core.utils.function_calling import convert_to_openai_tool
 from pydantic import ValidationError
+
 from ..langchain.utils import propagate_the_input_mapping
 
 logger = logging.getLogger(__name__)
 
 
+def replace_escaped_newlines(data):
+    """
+        Replace \\n with \n in all string values recursively.
+        Required for sanitization of state variables in code node
+    """
+    if isinstance(data, dict):
+        return {key: replace_escaped_newlines(value) for key, value in data.items()}
+    elif isinstance(data, str):
+        return data.replace('\\n', '\n')
+    else:
+        return data
+
 class FunctionTool(BaseTool):
     name: str = 'FunctionalTool'
     description: str = 'This is direct call node for tools'
@@ -21,6 +36,61 @@ class FunctionTool(BaseTool):
     input_variables: Optional[list[str]] = None
     input_mapping: Optional[dict[str, dict]] = None
     output_variables: Optional[list[str]] = None
+    structured_output: Optional[bool] = False
+    alita_client: Optional[Any] = None
+
+    def _prepare_pyodide_input(self, state: Union[str, dict, ToolCall]) -> str:
+        """Prepare input for PyodideSandboxTool by injecting state into the code block."""
+        # add state into the code block here since it might be changed during the execution of the code
+        state_copy = replace_escaped_newlines(deepcopy(state))
+
+        del state_copy['messages']  # remove messages to avoid issues with pickling without langchain-core
+        # inject state into the code block as alita_state variable
+        state_json = json.dumps(state_copy, ensure_ascii=False)
+        pyodide_predata = f'#state dict\nimport json\nalita_state = json.loads({json.dumps(state_json)})\n'
+
+        return pyodide_predata
+
+    def _handle_pyodide_output(self, tool_result: Any) -> dict:
+        """Handle output processing for PyodideSandboxTool results."""
+        tool_result_converted = {}
+
+        if self.output_variables:
+            for var in self.output_variables:
+                if var == "messages":
+                    tool_result_converted.update(
+                        {"messages": [{"role": "assistant", "content": dumps(tool_result)}]})
+                    continue
+                if isinstance(tool_result, dict) and var in tool_result:
+                    tool_result_converted[var] = tool_result[var]
+                else:
+                    # handler in case user points to a var that is not in the output of the tool
+                    tool_result_converted[var] = tool_result.get('result',
+                                                                 tool_result.get('error') if tool_result.get('error')
+                                                                 else 'Execution result is missing')
+        else:
+            tool_result_converted.update({"messages": [{"role": "assistant", "content": dumps(tool_result)}]})
+
+        if self.structured_output:
+            # execute code tool and update state variables
+            try:
+                result_value = tool_result.get('result', {})
+                if isinstance(result_value, dict):
+                    tool_result_converted.update(result_value)
+                elif isinstance(result_value, list):
+                    # Handle list case - could wrap in a key or handle differently based on requirements
+                    tool_result_converted.update({"result": result_value})
+                else:
+                    # Handle JSON string case
+                    tool_result_converted.update(json.loads(result_value))
+            except json.JSONDecodeError:
+                logger.error(f"JSONDecodeError: {tool_result}")
+
+        return tool_result_converted
+
+    def _is_pyodide_tool(self) -> bool:
+        """Check if the current tool is a PyodideSandboxTool."""
+        return self.tool.name.lower() == 'pyodide_sandbox'
 
     def invoke(
             self,
@@ -31,8 +101,15 @@ class FunctionTool(BaseTool):
         params = convert_to_openai_tool(self.tool).get(
             'function', {'parameters': {}}).get(
             'parameters', {'properties': {}}).get('properties', {})
+
         func_args = propagate_the_input_mapping(input_mapping=self.input_mapping, input_variables=self.input_variables,
                                                 state=state)
+
+        # special handler for PyodideSandboxTool
+        if self._is_pyodide_tool():
+            # replace new lines in strings in code block
+            code = func_args['code'].replace('\\n', '\\\\n')
+            func_args['code'] = f"{self._prepare_pyodide_input(state)}\n{code}"
         try:
             tool_result = self.tool.invoke(func_args, config, **kwargs)
             dispatch_custom_event(
@@ -44,11 +121,30 @@ class FunctionTool(BaseTool):
                 }, config=config
             )
             logger.info(f"ToolNode response: {tool_result}")
+
+            # handler for PyodideSandboxTool
+            if self._is_pyodide_tool():
+                return self._handle_pyodide_output(tool_result)
+
             if not self.output_variables:
                 return {"messages": [{"role": "assistant", "content": dumps(tool_result)}]}
             else:
-                if self.output_variables[0] == "messages":
-                    return {"messages": [{"role": "assistant", "content": dumps(tool_result)}]}
+                if "messages" in self.output_variables:
+                    messages_dict = {
+                        "messages": [{
+                            "role": "assistant",
+                            "content": dumps(tool_result)
+                            if not isinstance(tool_result, ToolException) and not isinstance(tool_result, str)
+                            else str(tool_result)
+                        }]
+                    }
+                    for var in self.output_variables:
+                        if var != "messages":
+                            if isinstance(tool_result, dict) and var in tool_result:
+                                messages_dict[var] = tool_result[var]
+                            else:
+                                messages_dict[var] = tool_result
+                    return messages_dict
                 else:
                     return { self.output_variables[0]: tool_result }
         except ValidationError:

alita_sdk/runtime/tools/graph.py

@@ -0,0 +1,81 @@
+import json
+
+from langgraph.graph.state import CompiledStateGraph
+
+from ..utils.utils import clean_string
+from langchain_core.tools import BaseTool
+from langchain_core.messages import BaseMessage, AIMessage, ToolCall
+from typing import Any, Type, Optional, Union
+from pydantic import create_model, field_validator, BaseModel
+from pydantic.fields import FieldInfo
+from ..langchain.mixedAgentRenderes import convert_message_to_json
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+graphToolSchema = create_model(
+    "graphToolSchema",
+    input=(str, FieldInfo(description="User Input for Graph")),
+    chat_history=(Optional[list[BaseMessage]],
+                  FieldInfo(description="Chat History relevant for Graph in format [{'role': '<user| assistant | etc>', 'content': '<content of the respected message>'}]", default=[]))
+)
+
+
+def formulate_query(kwargs):
+    chat_history = []
+    if kwargs.get('chat_history'):
+        if isinstance(kwargs.get('chat_history')[-1], BaseMessage):
+            chat_history = convert_message_to_json(kwargs.get('chat_history')[:])
+        elif isinstance(kwargs.get('chat_history')[-1], dict):
+            if all([True if message.get('role') and message.get('content') else False for message in
+                    kwargs.get('chat_history')]):
+                chat_history = kwargs.get('chat_history')[:]
+            else:
+                for each in kwargs.get('chat_history')[:]:
+                    chat_history.append(AIMessage(json.dumps(each)))
+        elif isinstance(kwargs.get('chat_history')[-1], str):
+            chat_history = []
+            for each in kwargs.get('chat_history')[:]:
+                chat_history.append(AIMessage(each))
+    elif kwargs.get('messages'):
+        chat_history = convert_message_to_json(kwargs.get('messages')[:])
+    result = {"input": kwargs.get('input'), "chat_history": chat_history}
+    for key, value in kwargs.items():
+        if key not in ("input", "chat_history"):
+            result[key] = value
+    return result
+
+
+class GraphTool(BaseTool):
+    name: str = 'GraphTool'
+    description: str = 'Graph tool for tools'
+    graph: CompiledStateGraph
+    args_schema: Type[BaseModel] = graphToolSchema
+    return_type: str = "str"
+
+    @field_validator('name', mode='before')
+    @classmethod
+    def remove_spaces(cls, v):
+        return clean_string(v)
+
+    def invoke(self, input: Any, config: Optional[dict] = None, **kwargs: Any) -> Any:
+        """Override default invoke to preserve all fields, not just args_schema"""
+        schema_values = self.args_schema(**input).model_dump() if self.args_schema else {}
+        extras = {k: v for k, v in input.items() if k not in schema_values}
+        all_kwargs = {**kwargs, **extras, **schema_values}
+        if config is None:
+            config = {}
+        # Pass the config to the _run empty or the one passed from the parent executor.
+        return self._run(config, **all_kwargs)
+
+    def _run(self, *args, **kwargs):
+        config = None
+        # From invoke method we are passing only 1 arg so it is safe to do this condition and config assignment.
+        # Default to None is safe because it will be checked also on the langchain side.
+        if args:
+            config = args[0]
+        response = self.graph.invoke(formulate_query(kwargs), config=config)
+        if self.return_type == "str":
+            return response["output"]
+        else:
+            return {"messages": [{"role": "assistant", "content": response["output"]}]}

alita_sdk/runtime/tools/image_generation.py

@@ -0,0 +1,218 @@
+"""
+Image generation tool for Alita SDK.
+"""
+import json
+import logging
+import uuid
+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"
+    )
+    n: int = Field(
+        default=1, description="Number of images to generate (1-10)",
+        ge=1, le=10
+    )
+    size: str = Field(
+        default="auto",
+        description="Size of the generated image (e.g., '1024x1024')"
+    )
+    quality: str = Field(
+        default="auto",
+        description="Quality of the generated image ('low', 'medium', 'high')"
+    )
+    style: Optional[str] = Field(
+        default=None, description="Style of the generated image (optional)"
+    )
+
+
+class ImageGenerationTool(BaseTool):
+    """Tool for generating images using the Alita client."""
+    
+    name: str = "generate_image"
+    description: str = (
+        "Generate images from text prompts using AI models. "
+        "Returns a JSON object with 'cached_image_id' field containing a reference to the generated image data. "
+        "The cached_image_id can be used to save or process the image. "
+        "The actual image data is stored temporarily and can be retrieved using the cached_image_id reference."
+    )
+    args_schema: Type[BaseModel] = ImageGenerationInput
+    alita_client: Any = None
+    
+    def __init__(self, client, **kwargs):
+        super().__init__(**kwargs)
+        self.alita_client = client
+    
+    def _run(self, prompt: str, n: int = 1, size: str = "auto",
+             quality: str = "auto", style: Optional[str] = None) -> str:
+        """Generate an image based on the provided parameters."""
+        try:
+            logger.debug(f"Generating image with prompt: {prompt[:50]}...")
+            
+            result = self.alita_client.generate_image(
+                prompt=prompt,
+                n=n,
+                size=size,
+                quality=quality,
+                style=style
+            )
+            
+            # Return simple JSON structure with reference ID instead of full base64
+            if 'data' in result:
+                images = result['data']
+                
+                # Process all images with unified structure
+                images_list = []
+                for idx, image_data in enumerate(images, 1):
+                    if not image_data.get('b64_json'):
+                        continue
+                    
+                    cached_image_id = f"img_{uuid.uuid4().hex[:12]}"
+                    
+                    # Store in cache
+                    if hasattr(self.alita_client, '_generated_images_cache'):
+                        self.alita_client._generated_images_cache[cached_image_id] = {
+                            'base64_data': image_data['b64_json']
+                        }
+                        logger.debug(f"Stored generated image in cache with ID: {cached_image_id}")
+                    
+                    images_list.append({
+                        "image_number": idx,
+                        "image_type": "png",
+                        "cached_image_id": cached_image_id
+                    })
+                
+                if not images_list:
+                    return json.dumps({
+                        "status": "error",
+                        "message": "No base64 image data found"
+                    })
+                
+                return json.dumps({
+                    "status": "success",
+                    "prompt": prompt,
+                    "total_images": len(images_list),
+                    "images": images_list
+                })
+            
+            # Fallback to error response if no images in result
+            return json.dumps({
+                "status": "error",
+                "message": f"Image generation completed but no images returned: {result}"
+            })
+            
+        except Exception as e:
+            logger.error(f"Error generating image: {e}")
+            return json.dumps({
+                "status": "error",
+                "message": f"Error generating image: {str(e)}"
+            })
+    
+    async def _arun(self, prompt: str, n: int = 1, size: str = "256x256",
+                    quality: str = "auto",
+                    style: Optional[str] = None) -> list:
+        """Async version - for now just calls the sync version."""
+        return self._run(prompt, n, size, quality, style)
+
+
+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