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

alita_sdk/runtime/langchain/document_loaders/constants.py

@@ -43,7 +43,7 @@ DEFAULT_ALLOWED_WITH_LLM = {
     LoaderProperties.PROMPT.value: "",
 }
 
-# Image file loaders mapping
+# Image file loaders mapping - directly supported by LLM with image_url
 image_loaders_map = {
     '.png': {
         'class': AlitaImageLoader,
@@ -73,6 +73,17 @@ image_loaders_map = {
         'kwargs': {},
         'allowed_to_override': DEFAULT_ALLOWED_WITH_LLM
     },
+    '.webp': {
+        'class': AlitaImageLoader,
+        'mime_type': 'image/webp',
+        'is_multimodal_processing': True,
+        'kwargs': {},
+        'allowed_to_override': DEFAULT_ALLOWED_WITH_LLM
+    }
+}
+
+# Image file loaders mapping - require conversion before sending to LLM
+image_loaders_map_converted = {
     '.bmp': {
         'class': AlitaImageLoader,
         'mime_type': 'image/bmp',
@@ -102,7 +113,7 @@ document_loaders_map = {
     },
     '.yml': {
         'class': AlitaTextLoader,
-        'mime_type': 'application/x-yaml',
+        'mime_type': 'application/yaml',
         'is_multimodal_processing': False,
         'kwargs': {
             'autodetect_encoding': True
@@ -111,7 +122,7 @@ document_loaders_map = {
     },
     '.yaml': {
         'class': AlitaTextLoader,
-        'mime_type': 'application/x-yaml',
+        'mime_type': 'application/yaml',
         'is_multimodal_processing': False,
         'kwargs': {
             'autodetect_encoding': True
@@ -244,17 +255,17 @@ document_loaders_map = {
             'extract_images': False,
         }
     },
-    '.py': {
-        'class': AlitaPythonLoader,
-        'mime_type': 'text/x-python',
-        'is_multimodal_processing': False,
-        'kwargs': {},
-        'allowed_to_override': DEFAULT_ALLOWED_BASE
-    }
+    # '.py': {
+    #     'class': AlitaPythonLoader,
+    #     'mime_type': 'text/x-python',
+    #     'is_multimodal_processing': False,
+    #     'kwargs': {},
+    #     'allowed_to_override': DEFAULT_ALLOWED_BASE
+    # }
 }
 
 code_extensions = [
-    # '.py',  # Python
+    '.py',  # Python
     '.js',  # JavaScript
     '.ts',  # TypeScript
     '.java',  # Java
@@ -292,7 +303,12 @@ default_loader_config = {
 code_loaders_map = {ext: default_loader_config for ext in code_extensions}
 
 # Combined mapping for backward compatibility
-loaders_map = {**image_loaders_map, **document_loaders_map, **code_loaders_map}
+loaders_map = {
+    **image_loaders_map,
+    **image_loaders_map_converted,
+    **document_loaders_map,
+    **code_loaders_map
+}
 
 loaders_allowed_to_override = {
     extension: config.get('allowed_to_override')

alita_sdk/runtime/langchain/langraph_agent.py

@@ -348,8 +348,8 @@ class StateModifierNode(Runnable):
         return result
 
 
-
-def prepare_output_schema(lg_builder, memory, store, debug=False, interrupt_before=None, interrupt_after=None, state_class=None, output_variables=None):
+def prepare_output_schema(lg_builder, memory, store, debug=False, interrupt_before=None, interrupt_after=None,
+                          state_class=None, output_variables=None):
     # prepare output channels
     if interrupt_after is None:
         interrupt_after = []
@@ -466,11 +466,12 @@ def create_graph(
                         elif node_type == 'agent':
                             input_params = node.get('input', ['messages'])
                             input_mapping = node.get('input_mapping',
-                                                       {'messages': {'type': 'variable', 'value': 'messages'}})
+                                                     {'messages': {'type': 'variable', 'value': 'messages'}})
+                            output_vars = node.get('output', [])
                             lg_builder.add_node(node_id, FunctionTool(
                                 client=client, tool=tool,
                                 name=node_id, return_type='str',
-                                output_variables=node.get('output', []),
+                                output_variables=output_vars + ['messages'] if 'messages' not in output_vars else output_vars,
                                 input_variables=input_params,
                                 input_mapping= input_mapping
                             ))
@@ -481,7 +482,8 @@ def create_graph(
                             # wrap with mappings
                             pipeline_name = node.get('tool', None)
                             if not pipeline_name:
-                                raise ValueError("Subgraph must have a 'tool' node: add required tool to the subgraph node")
+                                raise ValueError(
+                                    "Subgraph must have a 'tool' node: add required tool to the subgraph node")
                             node_fn = SubgraphRunnable(
                                 inner=tool.graph,
                                 name=pipeline_name,
@@ -499,15 +501,6 @@ def create_graph(
                                 structured_output=node.get('structured_output', False),
                                 task=node.get('task')
                             ))
-                        # TODO: decide on struct output for agent nodes
-                        # elif node_type == 'agent':
-                        #     lg_builder.add_node(node_id, AgentNode(
-                        #         client=client, tool=tool,
-                        #         name=node['id'], return_type='dict',
-                        #         output_variables=node.get('output', []),
-                        #         input_variables=node.get('input', ['messages']),
-                        #         task=node.get('task')
-                        #     ))
                         elif node_type == 'loop':
                             lg_builder.add_node(node_id, LoopNode(
                                 client=client, tool=tool,
@@ -520,7 +513,8 @@ def create_graph(
                             loop_toolkit_name = node.get('loop_toolkit_name')
                             loop_tool_name = node.get('loop_tool')
                             if (loop_toolkit_name and loop_tool_name) or loop_tool_name:
-                                loop_tool_name = f"{clean_string(loop_toolkit_name)}{TOOLKIT_SPLITTER}{loop_tool_name}" if loop_toolkit_name else clean_string(loop_tool_name)
+                                loop_tool_name = f"{clean_string(loop_toolkit_name)}{TOOLKIT_SPLITTER}{loop_tool_name}" if loop_toolkit_name else clean_string(
+                                    loop_tool_name)
                                 for t in tools:
                                     if t.name == loop_tool_name:
                                         logger.debug(f"Loop tool discovered: {t}")
@@ -555,12 +549,13 @@ def create_graph(
                         break
             elif node_type == 'code':
                 from ..tools.sandbox import create_sandbox_tool
-                sandbox_tool = create_sandbox_tool(stateful=False, allow_net=True)
-                code = node.get('code', "return 'Code block is empty'")
+                sandbox_tool = create_sandbox_tool(stateful=False, allow_net=True,
+                                                   alita_client=kwargs.get('alita_client', None))
+                code_data = node.get('code', {'type': 'fixed', 'value': "return 'Code block is empty'"})
                 lg_builder.add_node(node_id, FunctionTool(
                     tool=sandbox_tool, name=node['id'], return_type='dict',
                     output_variables=node.get('output', []),
-                    input_mapping={'code': {'type': 'fixed', 'value': code}},
+                    input_mapping={'code': code_data},
                     input_variables=node.get('input', ['messages']),
                     structured_output=node.get('structured_output', False),
                     alita_client=kwargs.get('alita_client', None)
@@ -593,7 +588,7 @@ def create_graph(
                 else:
                     # Use all available tools
                     available_tools = [tool for tool in tools if isinstance(tool, BaseTool)]
-                
+
                 lg_builder.add_node(node_id, LLMNode(
                     client=client,
                     input_mapping=node.get('input_mapping', {'messages': {'type': 'variable', 'value': 'messages'}}),
@@ -604,7 +599,9 @@ def create_graph(
                     input_variables=node.get('input', ['messages']),
                     structured_output=node.get('structured_output', False),
                     available_tools=available_tools,
-                    tool_names=tool_names))
+                    tool_names=tool_names,
+                    steps_limit=kwargs.get('steps_limit', 25)
+                ))
             elif node_type == 'router':
                 # Add a RouterNode as an independent node
                 lg_builder.add_node(node_id, RouterNode(
@@ -777,17 +774,46 @@ class LangGraphAgentRunnable(CompiledStateGraph):
             # Convert chat history dict messages to LangChain message objects
             chat_history = input.pop('chat_history')
             input['messages'] = [convert_dict_to_message(msg) for msg in chat_history]
-        
+
+        # handler for LLM node: if no input (Chat perspective), then take last human message
+        if not input.get('input'):
+            if input.get('messages'):
+                input['input'] = [next((msg for msg in reversed(input['messages']) if isinstance(msg, HumanMessage)),
+                                          None)]
+
         # Append current input to existing messages instead of overwriting
         if input.get('input'):
-            current_message = input.get('input')[-1]
+            if isinstance(input['input'], str):
+                current_message = input['input']
+            else:
+                current_message = input.get('input')[-1]
+
             # TODO: add handler after we add 2+ inputs (filterByType, etc.)
-            input['input'] = current_message if isinstance(current_message, str) else str(current_message)
+            if isinstance(current_message, HumanMessage):
+                current_content = current_message.content
+                if isinstance(current_content, list):
+                    text_contents = [
+                        item['text'] if isinstance(item, dict) and item.get('type') == 'text'
+                        else item if isinstance(item, str)
+                        else None
+                        for item in current_content
+                    ]
+                    text_contents = [text for text in text_contents if text is not None]
+                    input['input'] = ". ".join(text_contents)
+                elif isinstance(current_content, str):
+                    # on regenerate case
+                    input['input'] = current_content
+                else:
+                    input['input'] = str(current_content)
+            elif isinstance(current_message, str):
+                input['input'] = current_message
+            else:
+                input['input'] = str(current_message)
             if input.get('messages'):
                 # Ensure existing messages are LangChain objects
                 input['messages'] = [convert_dict_to_message(msg) for msg in input['messages']]
                 # Append to existing messages
-                input['messages'].append(current_message)
+                # input['messages'].append(current_message)
             else:
                 # No existing messages, create new list
                 input['messages'] = [current_message]
@@ -798,7 +824,8 @@ class LangGraphAgentRunnable(CompiledStateGraph):
         else:
             result = super().invoke(input, config=config, *args, **kwargs)
         try:
-            output = next((msg.content for msg in reversed(result['messages']) if not isinstance(msg, HumanMessage)), result['messages'][-1].content)
+            output = next((msg.content for msg in reversed(result['messages']) if not isinstance(msg, HumanMessage)),
+                          result['messages'][-1].content)
         except:
             output = list(result.values())[-1]
         config_state = self.get_state(config)
@@ -806,8 +833,6 @@ class LangGraphAgentRunnable(CompiledStateGraph):
         if is_execution_finished:
             thread_id = None
 
-
-
         result_with_state = {
             "output": output,
             "thread_id": thread_id,

alita_sdk/runtime/langchain/utils.py

@@ -5,8 +5,9 @@ import re
 from pydantic import create_model, Field
 from typing import Tuple, TypedDict, Any, Optional, Annotated
 from langchain_core.messages import AnyMessage
-from langchain_core.prompts import PromptTemplate
-from langgraph.graph import MessagesState, add_messages
+from langgraph.graph import add_messages
+
+from ...runtime.langchain.constants import ELITEA_RS
 
 logger = logging.getLogger(__name__)
 
@@ -130,7 +131,7 @@ def parse_type(type_str):
 
 
 def create_state(data: Optional[dict] = None):
-    state_dict = {'input': str, 'router_output': str}  # Always include router_output
+    state_dict = {'input': str, 'router_output': str, ELITEA_RS: str}  # Always include router_output
     types_dict = {}
     if not data:
         data = {'messages': 'list[str]'}
@@ -177,7 +178,17 @@ def propagate_the_input_mapping(input_mapping: dict[str, dict], input_variables:
             var_dict = create_params(input_variables, source)
 
         if value['type'] == 'fstring':
-            input_data[key] = value['value'].format(**var_dict)
+            try:
+                input_data[key] = value['value'].format(**var_dict)
+            except KeyError as e:
+                logger.error(f"KeyError in fstring formatting for key '{key}'. Attempt to find proper data in state.\n{e}")
+                try:
+                    # search for variables in state if not found in var_dict
+                    input_data[key] = value['value'].format(**state)
+                except KeyError as no_var_exception:
+                    logger.error(f"KeyError in fstring formatting for key '{key}' with state data.\n{no_var_exception}")
+                    # leave value as is if still not found (could be a constant string marked as fstring by mistake)
+                    input_data[key] = value['value']
         elif value['type'] == 'fixed':
             input_data[key] = value['value']
         else:

alita_sdk/runtime/toolkits/application.py

@@ -39,7 +39,14 @@ class ApplicationToolkit(BaseToolkit):
                                       description=app_details.get("description"), 
                                       application=app, 
                                       args_schema=applicationToolSchema,
-                                      return_type='str')])
+                                      return_type='str',
+                                      client=client,
+                                      args_runnable={
+                                          "application_id": application_id,
+                                          "application_version_id": application_version_id,
+                                          "store": store,
+                                          "llm": client.get_llm(version_details['llm_settings']['model_name'], model_settings),
+                                      })])
             
     def get_tools(self):
         return self.tools

alita_sdk/runtime/toolkits/tools.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional
 
 from langchain_core.tools import ToolException
 from langgraph.store.base import BaseStore
@@ -12,6 +13,8 @@ from .prompt import PromptToolkit
 from .subgraph import SubgraphToolkit
 from .vectorstore import VectorStoreToolkit
 from ..tools.mcp_server_tool import McpServerTool
+from ..tools.sandbox import SandboxToolkit
+from ..tools.image_generation import ImageGenerationToolkit
 # Import community tools
 from ...community import get_toolkits as community_toolkits, get_tools as community_tools
 from ...tools.memory import MemoryToolkit
@@ -24,64 +27,91 @@ def get_toolkits():
     core_toolkits = [
         ArtifactToolkit.toolkit_config_schema(),
         MemoryToolkit.toolkit_config_schema(),
-        VectorStoreToolkit.toolkit_config_schema()
+        VectorStoreToolkit.toolkit_config_schema(),
+        SandboxToolkit.toolkit_config_schema(),
+        ImageGenerationToolkit.toolkit_config_schema()
     ]
 
     return core_toolkits + community_toolkits() + alita_toolkits()
 
 
-def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None) -> list:
+def get_tools(tools_list: list, alita_client, llm, memory_store: BaseStore = None, debug_mode: Optional[bool] = False) -> list:
     prompts = []
     tools = []
 
     for tool in tools_list:
-        if tool['type'] == 'datasource':
-            tools.extend(DatasourcesToolkit.get_toolkit(
-                alita_client,
-                datasource_ids=[int(tool['settings']['datasource_id'])],
-                selected_tools=tool['settings']['selected_tools'],
-                toolkit_name=tool.get('toolkit_name', '') or tool.get('name', '')
-            ).get_tools())
-        elif tool['type'] == 'application' and tool.get('agent_type', '') != 'pipeline' :
-            tools.extend(ApplicationToolkit.get_toolkit(
-                alita_client,
-                application_id=int(tool['settings']['application_id']),
-                application_version_id=int(tool['settings']['application_version_id']),
-                selected_tools=[]
-            ).get_tools())
-        elif tool['type'] == 'application' and tool.get('agent_type', '') == 'pipeline':
-            # static get_toolkit returns a list of CompiledStateGraph stubs
-            tools.extend(SubgraphToolkit.get_toolkit(
-                alita_client,
-                application_id=int(tool['settings']['application_id']),
-                application_version_id=int(tool['settings']['application_version_id']),
-                app_api_key=alita_client.auth_token,
-                selected_tools=[],
-                llm=llm
-            ))
-        elif tool['type'] == 'memory':
-            tools += MemoryToolkit.get_toolkit(
-                namespace=tool['settings'].get('namespace', str(tool['id'])),
-                pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
-                store=memory_store,
-            ).get_tools()
-        elif tool['type'] == 'artifact':
-            tools.extend(ArtifactToolkit.get_toolkit(
-                client=alita_client,
-                bucket=tool['settings']['bucket'],
-                toolkit_name=tool.get('toolkit_name', ''),
-                selected_tools=tool['settings'].get('selected_tools', []),
-                llm=llm,
-                # indexer settings
-                pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
-                embedding_model=tool['settings'].get('embedding_model'),
-                collection_name=f"{tool.get('toolkit_name')}",
-            ).get_tools())
-        elif tool['type'] == 'vectorstore':
-            tools.extend(VectorStoreToolkit.get_toolkit(
-                llm=llm,
-                toolkit_name=tool.get('toolkit_name', ''),
-                **tool['settings']).get_tools())
+        try:
+            if tool['type'] == 'datasource':
+                tools.extend(DatasourcesToolkit.get_toolkit(
+                    alita_client,
+                    datasource_ids=[int(tool['settings']['datasource_id'])],
+                    selected_tools=tool['settings']['selected_tools'],
+                    toolkit_name=tool.get('toolkit_name', '') or tool.get('name', '')
+                ).get_tools())
+            elif tool['type'] == 'application' and tool.get('agent_type', '') != 'pipeline' :
+                tools.extend(ApplicationToolkit.get_toolkit(
+                    alita_client,
+                    application_id=int(tool['settings']['application_id']),
+                    application_version_id=int(tool['settings']['application_version_id']),
+                    selected_tools=[]
+                ).get_tools())
+            elif tool['type'] == 'application' and tool.get('agent_type', '') == 'pipeline':
+                # static get_toolkit returns a list of CompiledStateGraph stubs
+                tools.extend(SubgraphToolkit.get_toolkit(
+                    alita_client,
+                    application_id=int(tool['settings']['application_id']),
+                    application_version_id=int(tool['settings']['application_version_id']),
+                    app_api_key=alita_client.auth_token,
+                    selected_tools=[],
+                    llm=llm
+                ))
+            elif tool['type'] == 'memory':
+                tools += MemoryToolkit.get_toolkit(
+                    namespace=tool['settings'].get('namespace', str(tool['id'])),
+                    pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
+                    store=memory_store,
+                ).get_tools()
+            # TODO: update configuration of internal tools
+            elif tool['type'] == 'internal_tool':
+                if tool['name'] == 'pyodide':
+                    tools += SandboxToolkit.get_toolkit(
+                        stateful=False,
+                        allow_net=True,
+                        alita_client=alita_client,
+                    ).get_tools()
+                elif tool['name'] == 'image_generation':
+                    if alita_client and alita_client.model_image_generation:
+                        tools += ImageGenerationToolkit.get_toolkit(
+                            client=alita_client,
+                        ).get_tools()
+                    else:
+                        logger.warning("Image generation internal tool requested "
+                                       "but no image generation model configured")
+            elif tool['type'] == 'artifact':
+                tools.extend(ArtifactToolkit.get_toolkit(
+                    client=alita_client,
+                    bucket=tool['settings']['bucket'],
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    selected_tools=tool['settings'].get('selected_tools', []),
+                    llm=llm,
+                    # indexer settings
+                    pgvector_configuration=tool['settings'].get('pgvector_configuration', {}),
+                    embedding_model=tool['settings'].get('embedding_model'),
+                    collection_name=f"{tool.get('toolkit_name')}",
+                    collection_schema = str(tool['id'])
+                ).get_tools())
+            elif tool['type'] == 'vectorstore':
+                tools.extend(VectorStoreToolkit.get_toolkit(
+                    llm=llm,
+                    toolkit_name=tool.get('toolkit_name', ''),
+                    **tool['settings']).get_tools())
+        except Exception as e:
+            logger.error(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}", exc_info=True)
+            if debug_mode:
+                logger.info("Skipping tool initialization error due to debug mode.")
+                continue
+            else:
+                raise ToolException(f"Error initializing toolkit for tool '{tool.get('name', 'unknown')}': {e}")
     
     if len(prompts) > 0:
         tools += PromptToolkit.get_toolkit(alita_client, prompts).get_tools()

alita_sdk/runtime/tools/__init__.py

@@ -5,7 +5,11 @@ This module provides various tools that can be used within LangGraph agents.
 
 from .sandbox import PyodideSandboxTool, StatefulPyodideSandboxTool, create_sandbox_tool
 from .echo import EchoTool
-from .image_generation import ImageGenerationTool, create_image_generation_tool
+from .image_generation import (
+    ImageGenerationTool,
+    create_image_generation_tool,
+    ImageGenerationToolkit
+)
 
 __all__ = [
     "PyodideSandboxTool",
@@ -13,5 +17,6 @@ __all__ = [
     "create_sandbox_tool",
     "EchoTool",
     "ImageGenerationTool",
+    "ImageGenerationToolkit",
     "create_image_generation_tool"
-]
+]

alita_sdk/runtime/tools/application.py

@@ -50,6 +50,8 @@ class Application(BaseTool):
     application: Any
     args_schema: Type[BaseModel] = applicationToolSchema
     return_type: str = "str"
+    client: Any
+    args_runnable: dict = {}
     
     @field_validator('name', mode='before')
     @classmethod
@@ -66,6 +68,11 @@ class Application(BaseTool):
         return self._run(*config, **all_kwargs)
 
     def _run(self, *args, **kwargs):
+        if self.client and self.args_runnable:
+            # Recreate new LanggraphAgentRunnable in order to reflect the current input_mapping (it can be dynamic for pipelines).
+            # Actually, for pipelines agent toolkits LanggraphAgentRunnable is created (for LLMNode) before pipeline's schema parsing.
+            application_variables = {k: {"name": k, "value": v} for k, v in kwargs.items()}
+            self.application = self.client.application(**self.args_runnable, application_variables=application_variables)
         response = self.application.invoke(formulate_query(kwargs))
         if self.return_type == "str":
             return response["output"]

alita_sdk/runtime/tools/function.py

@@ -7,7 +7,7 @@ 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
 
@@ -31,24 +31,10 @@ class FunctionTool(BaseTool):
         """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)
-        # pickle state
-        import pickle
 
         del state_copy['messages']  # remove messages to avoid issues with pickling without langchain-core
-        serialized_state = pickle.dumps(state_copy)
         # inject state into the code block as alita_state variable
-        pyodide_predata = f"""import pickle\nalita_state = pickle.loads({serialized_state})\n"""
-        # add classes related to sandbox client
-        # read the content of alita_sdk/runtime/cliens/sandbox_client.py
-        try:
-            with open('alita_sdk/runtime/clients/sandbox_client.py', '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("sandbox_client.py not found. Ensure 'alita_sdk/runtime/clients/sandbox_client.py' exists.")
+        pyodide_predata = f"#state dict\nalita_state = {state_copy}\n"
         return pyodide_predata
 
     def _handle_pyodide_output(self, tool_result: Any) -> dict:
@@ -57,6 +43,10 @@ class FunctionTool(BaseTool):
 
         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:
@@ -71,8 +61,14 @@ class FunctionTool(BaseTool):
             # execute code tool and update state variables
             try:
                 result_value = tool_result.get('result', {})
-                tool_result_converted.update(result_value if isinstance(result_value, dict)
-                                           else json.loads(result_value))
+                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}")
 
@@ -98,7 +94,9 @@ class FunctionTool(BaseTool):
         # special handler for PyodideSandboxTool
         if self._is_pyodide_tool():
             code = func_args['code']
-            func_args['code'] = f"{self._prepare_pyodide_input(state)}\n{code}"
+            func_args['code'] = (f"{self._prepare_pyodide_input(state)}\n{code}"
+                                # handle new lines in the code properly
+                                 .replace('\\n','\\\\n'))
         try:
             tool_result = self.tool.invoke(func_args, config, **kwargs)
             dispatch_custom_event(
@@ -118,14 +116,21 @@ class FunctionTool(BaseTool):
             if not self.output_variables:
                 return {"messages": [{"role": "assistant", "content": dumps(tool_result)}]}
             else:
-                if self.output_variables[0] == "messages":
-                    return {
+                if "messages" in self.output_variables:
+                    messages_dict = {
                         "messages": [{
                             "role": "assistant",
-                            "content": dumps(tool_result) if not isinstance(tool_result, ToolException) else str(
-                                tool_result)
+                            "content": dumps(tool_result) if not isinstance(tool_result, ToolException)
+                            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: