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

alita_sdk/runtime/tools/data_analysis.py

@@ -0,0 +1,183 @@
+"""
+Data Analysis internal tool for Alita SDK.
+
+This tool provides Pandas-based data analysis capabilities as an internal tool,
+accessible through the "Enable internal tools" dropdown menu.
+
+It uses the conversation attachment bucket for file storage, providing seamless
+integration with drag-and-drop file uploads in chat.
+"""
+import logging
+from typing import Any, List, Literal, Optional, Type
+
+from langchain_core.tools import BaseTool, BaseToolkit
+from pydantic import BaseModel, ConfigDict, create_model, Field
+
+logger = logging.getLogger(__name__)
+
+name = "data_analysis"
+
+
+def get_tools(tools_list: list, alita_client=None, llm=None, memory_store=None):
+    """
+    Get data analysis tools for the provided tool configurations.
+
+    Args:
+        tools_list: List of tool configurations
+        alita_client: Alita client instance (required for data analysis)
+        llm: LLM client instance (required for code generation)
+        memory_store: Optional memory store instance (unused)
+
+    Returns:
+        List of data analysis tools
+    """
+    all_tools = []
+
+    for tool in tools_list:
+        if (tool.get('type') == 'data_analysis' or
+                tool.get('toolkit_name') == 'data_analysis'):
+            try:
+                if not alita_client:
+                    logger.error("Alita client is required for data analysis tools")
+                    continue
+
+                settings = tool.get('settings', {})
+                bucket_name = settings.get('bucket_name')
+
+                if not bucket_name:
+                    logger.error("bucket_name is required for data analysis tools")
+                    continue
+
+                toolkit_instance = DataAnalysisToolkit.get_toolkit(
+                    alita_client=alita_client,
+                    llm=llm,
+                    bucket_name=bucket_name,
+                    toolkit_name=tool.get('toolkit_name', '')
+                )
+                all_tools.extend(toolkit_instance.get_tools())
+            except Exception as e:
+                logger.error(f"Error in data analysis toolkit get_tools: {e}")
+                logger.error(f"Tool config: {tool}")
+                raise
+
+    return all_tools
+
+
+class DataAnalysisToolkit(BaseToolkit):
+    """
+    Data Analysis toolkit providing Pandas-based data analysis capabilities.
+
+    This is an internal tool that uses the conversation attachment bucket
+    for file storage, enabling seamless integration with chat file uploads.
+    """
+    tools: List[BaseTool] = []
+
+    @staticmethod
+    def toolkit_config_schema() -> Type[BaseModel]:
+        """Get the configuration schema for the data analysis toolkit."""
+        # Import PandasWrapper to get available tools schema
+        from alita_sdk.tools.pandas.api_wrapper import PandasWrapper
+
+        selected_tools = {
+            x['name']: x['args_schema'].model_json_schema()
+            for x in PandasWrapper.model_construct().get_available_tools()
+        }
+
+        return create_model(
+            'data_analysis',
+            bucket_name=(
+                Optional[str],
+                Field(
+                    default=None,
+                    title="Bucket name",
+                    description="Bucket where files are stored (auto-injected from conversation)"
+                )
+            ),
+            selected_tools=(
+                List[Literal[tuple(selected_tools)]],
+                Field(
+                    default=[],
+                    json_schema_extra={'args_schemas': selected_tools}
+                )
+            ),
+            __config__=ConfigDict(json_schema_extra={
+                'metadata': {
+                    "label": "Data Analysis",
+                    "icon_url": "data-analysis.svg",
+                    "hidden": True,  # Hidden from regular toolkit menu
+                    "categories": ["internal_tool"],
+                    "extra_categories": ["data analysis", "pandas", "dataframes", "data science"],
+                }
+            })
+        )
+
+    @classmethod
+    def get_toolkit(
+        cls,
+        alita_client=None,
+        llm=None,
+        bucket_name: str = None,
+        toolkit_name: Optional[str] = None,
+        selected_tools: Optional[List[str]] = None,
+        **kwargs
+    ):
+        """
+        Get toolkit with data analysis tools.
+
+        Args:
+            alita_client: Alita client instance (required)
+            llm: LLM for code generation (optional, uses alita_client.llm if not provided)
+            bucket_name: Conversation attachment bucket (required)
+            toolkit_name: Optional name prefix for tools
+            selected_tools: Optional list of tool names to include (default: all)
+            **kwargs: Additional arguments
+
+        Returns:
+            DataAnalysisToolkit instance with configured tools
+
+        Raises:
+            ValueError: If alita_client or bucket_name is not provided
+        """
+        if not alita_client:
+            raise ValueError("Alita client is required for data analysis")
+
+        if not bucket_name:
+            raise ValueError("bucket_name is required for data analysis (should be conversation attachment bucket)")
+
+        # Import the PandasWrapper from existing toolkit
+        from alita_sdk.tools.pandas.api_wrapper import PandasWrapper
+        from alita_sdk.tools.base.tool import BaseAction
+
+        # Create wrapper with conversation bucket
+        wrapper = PandasWrapper(
+            alita=alita_client,
+            llm=llm,
+            bucket_name=bucket_name
+        )
+
+        # Get tools from wrapper
+        available_tools = wrapper.get_available_tools()
+        tools = []
+
+        for tool in available_tools:
+            # Filter by selected_tools if provided
+            if selected_tools and tool["name"] not in selected_tools:
+                continue
+
+            description = tool["description"]
+            if toolkit_name:
+                description = f"Toolkit: {toolkit_name}\n{description}"
+            description = description[:1000]
+
+            tools.append(BaseAction(
+                api_wrapper=wrapper,
+                name=tool["name"],
+                description=description,
+                args_schema=tool["args_schema"],
+                metadata={"toolkit_name": toolkit_name, "toolkit_type": name} if toolkit_name else {}
+            ))
+
+        return cls(tools=tools)
+
+    def get_tools(self):
+        return self.tools

alita_sdk/runtime/tools/function.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import logging
 from copy import deepcopy
@@ -7,15 +8,28 @@ 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, ToolException
-from typing import Any, Optional, Union, Annotated
+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
+from ..utils.serialization import safe_serialize
 
 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'
@@ -28,31 +42,36 @@ class FunctionTool(BaseTool):
     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 = 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
-        pyodide_predata = f"alita_state = {state_copy}"
-        # add classes related to sandbox client
-        # read the content of alita_sdk/runtime/cliens/sandbox_client.py
-        try:
-            import os
-            from pathlib import Path
-
-            # Get the directory of the current file and construct the path to sandbox_client.py
-            current_dir = Path(__file__).parent
-            sandbox_client_path = current_dir.parent / 'clients' / 'sandbox_client.py'
-
-            with open(sandbox_client_path, 'r') as f:
-                sandbox_client_code = f.read()
-            pyodide_predata += f"\n{sandbox_client_code}\n"
-            pyodide_predata += (f"alita_client = SandboxClient(base_url='{self.alita_client.base_url}',"
-                                f"project_id={self.alita_client.project_id},"
-                                f"auth_token='{self.alita_client.auth_token}')")
-        except FileNotFoundError:
-            logger.error(f"sandbox_client.py not found at {sandbox_client_path}. Ensure the file exists.")
+        """Prepare input for PyodideSandboxTool by injecting state into the code block.
+
+        Uses base64 encoding to avoid string escaping issues when passing JSON
+        through multiple layers of parsing (Python -> Deno -> Pyodide) and compression to minimize args list
+        """
+        import base64
+        import zlib
+
+        state_copy = replace_escaped_newlines(deepcopy(state))
+        if 'messages' in state_copy:
+            del state_copy['messages']
+
+        # Use safe_serialize to handle Pydantic models, datetime, and other non-JSON types
+        state_json = safe_serialize(state_copy)
+
+        # Use base64 encoding to avoid all string escaping issues
+        # This is more robust than repr() when the code passes through multiple parsers
+        # use compression to avoid issue with `{"error": "Error executing code: [Errno 7] Argument list too long: 'deno'"}`
+        compressed = zlib.compress(state_json.encode('utf-8'))
+        encoded = base64.b64encode(compressed).decode('ascii')
+
+        pyodide_predata = f'''#state dict
+import json
+import base64
+import zlib
+
+compressed_state = base64.b64decode('{encoded}')
+state_json = zlib.decompress(compressed_state).decode('utf-8')
+alita_state = json.loads(state_json)
+'''
         return pyodide_predata
 
     def _handle_pyodide_output(self, tool_result: Any) -> dict:
@@ -111,7 +130,8 @@ class FunctionTool(BaseTool):
 
         # special handler for PyodideSandboxTool
         if self._is_pyodide_tool():
-            code = func_args['code']
+            # 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)
@@ -132,14 +152,26 @@ class FunctionTool(BaseTool):
             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 not isinstance(tool_result, ToolException) else str(
-                                tool_result)
-                        }]
-                    }
+                if "messages" in self.output_variables:
+                    if 'messages' in tool_result:
+                        # case when the sub-graph has been executed
+                        messages_dict = {"messages": tool_result['messages']}
+                    else:
+                        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

@@ -47,8 +47,8 @@ def formulate_query(kwargs):
 
 
 class GraphTool(BaseTool):
-    name: str
-    description: str
+    name: str = 'GraphTool'
+    description: str = 'Graph tool for tools'
     graph: CompiledStateGraph
     args_schema: Type[BaseModel] = graphToolSchema
     return_type: str = "str"
@@ -65,10 +65,16 @@ class GraphTool(BaseTool):
         all_kwargs = {**kwargs, **extras, **schema_values}
         if config is None:
             config = {}
-        return self._run(*config, **all_kwargs)
+        # 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):
-        response = self.graph.invoke(formulate_query(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:

alita_sdk/runtime/tools/image_generation.py

@@ -1,17 +1,62 @@
 """
 Image generation tool for Alita SDK.
 """
+import json
 import logging
-from typing import Optional, Type, Any
-from langchain_core.tools import BaseTool
-from pydantic import BaseModel, Field
+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")
+    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 +67,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)"
@@ -33,7 +78,12 @@ 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"
+    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
     
@@ -42,10 +92,10 @@ class ImageGenerationTool(BaseTool):
         self.alita_client = client
     
     def _run(self, prompt: str, n: int = 1, size: str = "auto",
-             quality: str = "auto", style: Optional[str] = None) -> list:
+             quality: str = "auto", style: Optional[str] = None) -> str:
         """Generate an image based on the provided parameters."""
         try:
-            logger.info(f"Generating image with prompt: {prompt[:50]}...")
+            logger.debug(f"Generating image with prompt: {prompt[:50]}...")
             
             result = self.alita_client.generate_image(
                 prompt=prompt,
@@ -55,54 +105,56 @@ class ImageGenerationTool(BaseTool):
                 style=style
             )
             
-            # Return multimodal content format for LLM consumption
+            # Return simple JSON structure with reference ID instead of full base64
             if 'data' in result:
                 images = result['data']
-                content_chunks = []
                 
-                # Add a text description of what was generated
-                if len(images) == 1:
-                    content_chunks.append({
-                        "type": "text",
-                        "text": f"Generated image for prompt: '{prompt}'"
-                    })
-                else:
-                    content_chunks.append({
-                        "type": "text",
-                        "text": f"Generated {len(images)} images for prompt: '{prompt}'"
+                # 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
                     })
                 
-                # Add image content for each generated image
-                for image_data in images:
-                    if image_data.get('url'):
-                        content_chunks.append({
-                            "type": "image_url",
-                            "image_url": {
-                                "url": image_data['url']
-                            }
-                        })
-                    elif image_data.get('b64_json'):
-                        content_chunks.append({
-                            "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/png;base64,{image_data['b64_json']}"
-                            }
-                        })
+                if not images_list:
+                    return json.dumps({
+                        "status": "error",
+                        "message": "No base64 image data found"
+                    })
                 
-                return content_chunks
+                return json.dumps({
+                    "status": "success",
+                    "prompt": prompt,
+                    "total_images": len(images_list),
+                    "images": images_list
+                })
             
-            # Fallback to text response if no images in result
-            return [{
-                "type": "text",
-                "text": f"Image generation completed but no images returned: {result}"
-            }]
+            # 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 [{
-                "type": "text",
-                "text": f"Error generating image: {str(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",
@@ -114,3 +166,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