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

alita_sdk/runtime/langchain/langraph_agent.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import Union, Any, Optional, Annotated, get_type_hints
 from uuid import uuid4
 from typing import Dict
@@ -11,6 +12,7 @@ from langchain_core.runnables import Runnable
 from langchain_core.runnables import RunnableConfig
 from langchain_core.tools import BaseTool, ToolException
 from langgraph.channels.ephemeral_value import EphemeralValue
+from langgraph.errors import GraphRecursionError
 from langgraph.graph import StateGraph
 from langgraph.graph.graph import END, START
 from langgraph.graph.state import CompiledStateGraph
@@ -18,8 +20,10 @@ from langgraph.managed.base import is_managed_value
 from langgraph.prebuilt import InjectedStore
 from langgraph.store.base import BaseStore
 
+from .constants import PRINTER_NODE_RS, PRINTER, PRINTER_COMPLETED_STATE
 from .mixedAgentRenderes import convert_message_to_json
-from .utils import create_state, propagate_the_input_mapping
+from .utils import create_state, propagate_the_input_mapping, safe_format
+from ..utils.constants import TOOLKIT_NAME_META, TOOL_NAME_META
 from ..tools.function import FunctionTool
 from ..tools.indexer_tool import IndexerNode
 from ..tools.llm import LLMNode
@@ -27,7 +31,7 @@ from ..tools.loop import LoopNode
 from ..tools.loop_output import LoopToolNode
 from ..tools.tool import ToolNode
 from ..utils.evaluate import EvaluateTemplate
-from ..utils.utils import clean_string, TOOLKIT_SPLITTER
+from ..utils.utils import clean_string
 from ..tools.router import RouterNode
 
 logger = logging.getLogger(__name__)
@@ -169,12 +173,13 @@ Answer only with step name, no need to add descrip in case none of the steps are
 """
 
     def __init__(self, client, steps: str, description: str = "", decisional_inputs: Optional[list[str]] = [],
-                 default_output: str = 'END'):
+                 default_output: str = 'END', is_node: bool = False):
         self.client = client
         self.steps = ",".join([clean_string(step) for step in steps])
         self.description = description
         self.decisional_inputs = decisional_inputs
         self.default_output = default_output if default_output != 'END' else END
+        self.is_node = is_node
 
     def invoke(self, state: Annotated[BaseStore, InjectedStore()], config: Optional[RunnableConfig] = None) -> str:
         additional_info = ""
@@ -184,10 +189,10 @@ Answer only with step name, no need to add descrip in case none of the steps are
                 decision_input = state.get('messages', [])[:]
             else:
                 if len(additional_info) == 0:
-                    additional_info = """### Additoinal info: """
+                    additional_info = """### Additional info: """
                 additional_info += "{field}: {value}\n".format(field=field, value=state.get(field, ""))
         decision_input.append(HumanMessage(
-            self.prompt.format(steps=self.steps, description=self.description, additional_info=additional_info)))
+            self.prompt.format(steps=self.steps, description=safe_format(self.description, state), additional_info=additional_info)))
         completion = self.client.invoke(decision_input)
         result = clean_string(completion.content.strip())
         logger.info(f"Plan to transition to: {result}")
@@ -196,7 +201,8 @@ Answer only with step name, no need to add descrip in case none of the steps are
         dispatch_custom_event(
             "on_decision_edge", {"decisional_inputs": self.decisional_inputs, "state": state}, config=config
         )
-        return result
+        # support of legacy `decision` as part of node
+        return {"router_output": result} if self.is_node else result
 
 
 class TransitionalEdge(Runnable):
@@ -231,6 +237,35 @@ class StateDefaultNode(Runnable):
                     result[key] = temp_value
         return result
 
+class PrinterNode(Runnable):
+    name = "PrinterNode"
+
+    def __init__(self, input_mapping: Optional[dict[str, dict]]):
+        self.input_mapping = input_mapping
+
+    def invoke(self, state: BaseStore, config: Optional[RunnableConfig] = None) -> dict:
+        logger.info(f"Printer Node - Current state variables: {state}")
+        result = {}
+        logger.debug(f"Initial text pattern: {self.input_mapping}")
+        mapping = propagate_the_input_mapping(self.input_mapping, [], state)
+        # for printer node we expect that all the lists will be joined into strings already
+        # Join any lists that haven't been converted yet
+        for key, value in mapping.items():
+            if isinstance(value, list):
+                mapping[key] = ', '.join(str(item) for item in value)
+        if mapping.get(PRINTER) is None:
+            raise ToolException(f"PrinterNode requires '{PRINTER}' field in input mapping")
+        formatted_output = mapping[PRINTER]
+        # add info label to the printer's output
+        if not formatted_output == PRINTER_COMPLETED_STATE:
+            # convert formatted output to string if it's not
+            if not isinstance(formatted_output, str):
+                formatted_output = str(formatted_output)
+            formatted_output += f"\n\n-----\n*How to proceed?*\n* *to resume the pipeline - type anything...*"
+        logger.debug(f"Formatted output: {formatted_output}")
+        result[PRINTER_NODE_RS] = formatted_output
+        return result
+
 
 class StateModifierNode(Runnable):
     name = "StateModifierNode"
@@ -248,19 +283,82 @@ class StateModifierNode(Runnable):
 
         # Collect input variables from state
         input_data = {}
+        
         for var in self.input_variables:
             if var in state:
                 input_data[var] = state.get(var)
-
+        type_of_output = type(state.get(self.output_variables[0])) if self.output_variables else None
         # Render the template using Jinja
-        from jinja2 import Template
-        rendered_message = Template(self.template).render(**input_data)
+        import json
+        import base64
+        from jinja2 import Environment
+
+        def from_json(value):
+            """Convert JSON string to Python object"""
+            try:
+                return json.loads(value)
+            except (json.JSONDecodeError, TypeError) as e:
+                logger.warning(f"Failed to parse JSON value: {e}")
+                return value
+        
+        def base64_to_string(value):
+            """Convert base64 encoded string to regular string"""
+            try:
+                return base64.b64decode(value).decode('utf-8')
+            except Exception as e:
+                logger.warning(f"Failed to decode base64 value: {e}")
+                return value
+        
+        def split_by_words(value, chunk_size=100):
+            words = value.split()
+            return [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
+        
+        def split_by_regex(value, pattern):
+            """Splits the provided string using the specified regex pattern."""
+            return re.split(pattern, value)
+
+        env = Environment()
+        env.filters['from_json'] = from_json
+        env.filters['base64_to_string'] = base64_to_string
+        env.filters['split_by_words'] = split_by_words
+        env.filters['split_by_regex'] = split_by_regex
+
+        template = env.from_string(self.template)
+        rendered_message = template.render(**input_data)
         result = {}
         # Store the rendered message in the state or messages
         if len(self.output_variables) > 0:
             # Use the first output variable to store the rendered content
             output_var = self.output_variables[0]
-            result[output_var] = rendered_message
+            
+            # Convert rendered_message to the appropriate type
+            if type_of_output is not None:
+                try:
+                    if type_of_output == dict:
+                        result[output_var] = json.loads(rendered_message) if isinstance(rendered_message, str) else dict(rendered_message)
+                    elif type_of_output == list:
+                        result[output_var] = json.loads(rendered_message) if isinstance(rendered_message, str) else list(rendered_message)
+                    elif type_of_output == int:
+                        result[output_var] = int(rendered_message)
+                    elif type_of_output == float:
+                        result[output_var] = float(rendered_message)
+                    elif type_of_output == str:
+                        result[output_var] = str(rendered_message)
+                    elif type_of_output == bool:
+                        if isinstance(rendered_message, str):
+                            result[output_var] = rendered_message.lower() in ('true', '1', 'yes', 'on')
+                        else:
+                            result[output_var] = bool(rendered_message)
+                    elif type_of_output == type(None):
+                        result[output_var] = None
+                    else:
+                        # Fallback to string if type is not recognized
+                        result[output_var] = str(rendered_message)
+                except (ValueError, TypeError, json.JSONDecodeError) as e:
+                    logger.warning(f"Failed to convert rendered_message to {type_of_output.__name__}: {e}. Using string fallback.")
+                    result[output_var] = str(rendered_message)
+            else:
+                result[output_var] = rendered_message
 
         # Clean up specified variables (make them empty, not delete)
         
@@ -284,8 +382,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 = []
@@ -350,6 +448,50 @@ def prepare_output_schema(lg_builder, memory, store, debug=False, interrupt_befo
     return compiled
 
 
+def find_tool_by_name_or_metadata(tools: list, tool_name: str, toolkit_name: Optional[str] = None) -> Optional[BaseTool]:
+    """
+    Find a tool by name or by matching metadata (toolkit_name + tool_name).
+
+    For toolkit nodes with toolkit_name specified, this function checks:
+    1. Metadata match first (toolkit_name + tool_name) - PRIORITY when toolkit_name is provided
+    2. Direct tool name match (backward compatibility fallback)
+
+    For toolkit nodes without toolkit_name, or other node types:
+    1. Direct tool name match
+
+    Args:
+        tools: List of available tools
+        tool_name: The tool name to search for
+        toolkit_name: Optional toolkit name for metadata matching
+
+    Returns:
+        The matching tool or None if not found
+    """
+    # When toolkit_name is specified, prioritize metadata matching
+    if toolkit_name:
+        for tool in tools:
+            # Check metadata match first
+            if hasattr(tool, 'metadata') and tool.metadata:
+                metadata_toolkit_name = tool.metadata.get(TOOLKIT_NAME_META)
+                metadata_tool_name = tool.metadata.get(TOOL_NAME_META)
+
+                # Match if both toolkit_name and tool_name in metadata match
+                if metadata_toolkit_name == toolkit_name and metadata_tool_name == tool_name:
+                    return tool
+
+        # Fallback to direct name match for backward compatibility
+        for tool in tools:
+            if tool.name == tool_name:
+                return tool
+    else:
+        # No toolkit_name specified, use direct name match only
+        for tool in tools:
+            if tool.name == tool_name:
+                return tool
+
+    return None
+
+
 def create_graph(
         client: Any,
         yaml_schema: str,
@@ -385,44 +527,64 @@ def create_graph(
             node_type = node.get('type', 'function')
             node_id = clean_string(node['id'])
             toolkit_name = node.get('toolkit_name')
-            tool_name = clean_string(node.get('tool', node_id))
-            if toolkit_name:
-                tool_name = f"{clean_string(toolkit_name)}{TOOLKIT_SPLITTER}{tool_name}"
+            tool_name = clean_string(node.get('tool', ''))
+            # Tool names are now clean (no prefix needed)
             logger.info(f"Node: {node_id} : {node_type} - {tool_name}")
-            if node_type in ['function', 'tool', 'loop', 'loop_from_tool', 'indexer', 'subgraph', 'pipeline', 'agent']:
-                for tool in tools:
-                    if tool.name == tool_name:
-                        if node_type == 'function':
+            if node_type in ['function', 'toolkit', 'mcp', 'tool', 'loop', 'loop_from_tool', 'indexer', 'subgraph', 'pipeline', 'agent']:
+                if node_type in ['mcp', 'toolkit', 'agent'] and not tool_name:
+                    # tool is not specified
+                    raise ToolException(f"Tool name is required for {node_type} node with id '{node_id}'")
+
+                # Unified validation and tool finding for toolkit, mcp, and agent node types
+                matching_tool = None
+                if node_type in ['toolkit', 'mcp', 'agent']:
+                    # Use enhanced validation that checks both direct name and metadata
+                    matching_tool = find_tool_by_name_or_metadata(tools, tool_name, toolkit_name)
+                    if not matching_tool:
+                        # tool is not found in the provided tools
+                        error_msg = f"Node `{node_id}` with type `{node_type}` has tool '{tool_name}'"
+                        if toolkit_name:
+                            error_msg += f" (toolkit: '{toolkit_name}')"
+                        error_msg += f" which is not found in the provided tools. Make sure it is connected properly. Available tools: {format_tools(tools)}"
+                        raise ToolException(error_msg)
+                else:
+                    # For other node types, find tool by direct name match
+                    for tool in tools:
+                        if tool.name == tool_name:
+                            matching_tool = tool
+                            break
+
+                if matching_tool:
+                        if node_type in ['function', 'toolkit', 'mcp']:
                             lg_builder.add_node(node_id, FunctionTool(
-                                tool=tool, name=node['id'], return_type='dict',
+                                tool=matching_tool, name=node_id, return_type='dict',
                                 output_variables=node.get('output', []),
                                 input_mapping=node.get('input_mapping',
                                                        {'messages': {'type': 'variable', 'value': 'messages'}}),
                                 input_variables=node.get('input', ['messages'])))
                         elif node_type == 'agent':
                             input_params = node.get('input', ['messages'])
-                            input_mapping = {'task': {'type': 'fstring', 'value': f"{node.get('task', '')}"},
-                                                      'chat_history': {'type': 'fixed', 'value': []}}
-                            # Add 'chat_history' to input_mapping only if 'messages' is in input_params
-                            if 'messages' in input_params:
-                                input_mapping['chat_history'] = {'type': 'variable', 'value': 'messages'}
+                            input_mapping = node.get('input_mapping',
+                                                     {'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='dict',
-                                output_variables=node.get('output', []),
+                                client=client, tool=matching_tool,
+                                name=node_id, return_type='str',
+                                output_variables=output_vars + ['messages'] if 'messages' not in output_vars else output_vars,
                                 input_variables=input_params,
                                 input_mapping= input_mapping
                             ))
                         elif node_type == 'subgraph' or node_type == 'pipeline':
                             # assign parent memory/store
-                            # tool.checkpointer = memory
-                            # tool.store = store
+                            # matching_tool.checkpointer = memory
+                            # matching_tool.store = store
                             # 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,
+                                inner=matching_tool.graph,
                                 name=pipeline_name,
                                 input_mapping=node.get('input_mapping', {}),
                                 output_mapping=node.get('output_mapping', {}),
@@ -431,26 +593,17 @@ def create_graph(
                             break  # skip legacy handling
                         elif node_type == 'tool':
                             lg_builder.add_node(node_id, ToolNode(
-                                client=client, tool=tool,
-                                name=node['id'], return_type='dict',
+                                client=client, tool=matching_tool,
+                                name=node_id, return_type='dict',
                                 output_variables=node.get('output', []),
                                 input_variables=node.get('input', ['messages']),
                                 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,
-                                name=node['id'], return_type='dict',
+                                client=client, tool=matching_tool,
+                                name=node_id, return_type='dict',
                                 output_variables=node.get('output', []),
                                 input_variables=node.get('input', ['messages']),
                                 task=node.get('task', '')
@@ -459,14 +612,15 @@ 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)
+                                # Use clean tool name (no prefix)
+                                loop_tool_name = clean_string(loop_tool_name)
                                 for t in tools:
                                     if t.name == loop_tool_name:
                                         logger.debug(f"Loop tool discovered: {t}")
                                         lg_builder.add_node(node_id, LoopToolNode(
                                             client=client,
-                                            name=node['id'], return_type='dict',
-                                            tool=tool, loop_tool=t,
+                                            name=node_id, return_type='dict',
+                                            tool=matching_tool, loop_tool=t,
                                             variables_mapping=node.get('variables_mapping', {}),
                                             output_variables=node.get('output', []),
                                             input_variables=node.get('input', ['messages']),
@@ -482,16 +636,28 @@ def create_graph(
                                     indexer_tool = t
                             logger.info(f"Indexer tool: {indexer_tool}")
                             lg_builder.add_node(node_id, IndexerNode(
-                                client=client, tool=tool,
+                                client=client, tool=matching_tool,
                                 index_tool=indexer_tool,
                                 input_mapping=node.get('input_mapping', {}),
-                                name=node['id'], return_type='dict',
+                                name=node_id, return_type='dict',
                                 chunking_tool=node.get('chunking_tool', None),
                                 chunking_config=node.get('chunking_config', {}),
                                 output_variables=node.get('output', []),
                                 input_variables=node.get('input', ['messages']),
                                 structured_output=node.get('structured_output', False)))
-                        break
+            elif node_type == 'code':
+                from ..tools.sandbox import create_sandbox_tool
+                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': code_data},
+                    input_variables=node.get('input', ['messages']),
+                    structured_output=node.get('structured_output', False),
+                    alita_client=kwargs.get('alita_client', None)
+                ))
             elif node_type == 'llm':
                 output_vars = node.get('output', [])
                 output_vars_dict = {
@@ -504,10 +670,10 @@ def create_graph(
                 tool_names = []
                 if isinstance(connected_tools, dict):
                     for toolkit, selected_tools in connected_tools.items():
-                        for tool in selected_tools:
-                            tool_names.append(f"{toolkit}{TOOLKIT_SPLITTER}{tool}")
+                        # Add tool names directly (no prefix)
+                        tool_names.extend(selected_tools)
                 elif isinstance(connected_tools, list):
-                    # for cases when tools are provided as a list of names with already bound toolkit_name
+                    # Use provided tool names as-is
                     tool_names = connected_tools
                 
                 if tool_names:
@@ -520,28 +686,44 @@ 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, 
-                    prompt=node.get('prompt', {}),
-                    name=node['id'], 
+                    client=client,
+                    input_mapping=node.get('input_mapping', {'messages': {'type': 'variable', 'value': 'messages'}}),
+                    name=node_id,
                     return_type='dict',
-                    response_key=node.get('response_key', 'messages'),
                     structured_output_dict=output_vars_dict,
                     output_variables=output_vars,
                     input_variables=node.get('input', ['messages']),
                     structured_output=node.get('structured_output', False),
+                    tool_execution_timeout=node.get('tool_execution_timeout', 900),
                     available_tools=available_tools,
-                    tool_names=tool_names))
-            elif node_type == 'router':
-                # Add a RouterNode as an independent node
-                lg_builder.add_node(node_id, RouterNode(
-                    name=node['id'],
-                    condition=node.get('condition', ''),
-                    routes=node.get('routes', []),
-                    default_output=node.get('default_output', 'END'),
-                    input_variables=node.get('input', ['messages'])
+                    tool_names=tool_names,
+                    steps_limit=kwargs.get('steps_limit', 25)
                 ))
+            elif node_type in ['router', 'decision']:
+                if node_type == 'router':
+                    # Add a RouterNode as an independent node
+                    lg_builder.add_node(node_id, RouterNode(
+                        name=node_id,
+                        condition=node.get('condition', ''),
+                        routes=node.get('routes', []),
+                        default_output=node.get('default_output', 'END'),
+                        input_variables=node.get('input', ['messages'])
+                    ))
+                elif node_type == 'decision':
+                    logger.info(f'Adding decision: {node["nodes"]}')
+                    # fallback to old-style decision node
+                    decisional_inputs = node.get('decisional_inputs')
+                    decisional_inputs = node.get('input', ['messages']) if not decisional_inputs else decisional_inputs
+                    lg_builder.add_node(node_id, DecisionEdge(
+                        client, node['nodes'],
+                        node.get('description', ""),
+                        decisional_inputs=decisional_inputs,
+                        default_output=node.get('default_output', 'END'),
+                        is_node=True
+                    ))
+
                 # Add a single conditional edge for all routes
                 lg_builder.add_conditional_edges(
                     node_id,
@@ -552,6 +734,7 @@ def create_graph(
                         default_output=node.get('default_output', 'END')
                     )
                 )
+                continue
             elif node_type == 'state_modifier':
                 lg_builder.add_node(node_id, StateModifierNode(
                     template=node.get('template', ''),
@@ -559,6 +742,22 @@ def create_graph(
                     input_variables=node.get('input', ['messages']),
                     output_variables=node.get('output', [])
                 ))
+            elif node_type == 'printer':
+                lg_builder.add_node(node_id, PrinterNode(
+                    input_mapping=node.get('input_mapping', {'printer': {'type': 'fixed', 'value': ''}}),
+                ))
+
+                # add interrupts after printer node if specified
+                interrupt_after.append(clean_string(node_id))
+
+                # reset printer output variable to avoid carrying over
+                reset_node_id = f"{node_id}_reset"
+                lg_builder.add_node(reset_node_id, PrinterNode(
+                    input_mapping={'printer': {'type': 'fixed', 'value': PRINTER_COMPLETED_STATE}}
+                ))
+                lg_builder.add_conditional_edges(node_id, TransitionalEdge(reset_node_id))
+                lg_builder.add_conditional_edges(reset_node_id, TransitionalEdge(clean_string(node['transition'])))
+                continue
             if node.get('transition'):
                 next_step = clean_string(node['transition'])
                 logger.info(f'Adding transition: {next_step}')
@@ -584,14 +783,11 @@ def create_graph(
             entry_point = clean_string(schema['entry_point'])
         except KeyError:
             raise ToolException("Entry point is not defined in the schema. Please define 'entry_point' in the schema.")
-        for key, value in state.items():
-            if 'type' in value and 'value' in value:
-                # set default value for state variable if it is defined in the schema
-                state_default_node = StateDefaultNode(default_vars=state)
-                lg_builder.add_node(state_default_node.name, state_default_node)
-                lg_builder.set_entry_point(state_default_node.name)
-                lg_builder.add_conditional_edges(state_default_node.name, TransitionalEdge(entry_point))
-                break
+        if state.items():
+            state_default_node = StateDefaultNode(default_vars=set_defaults(state))
+            lg_builder.add_node(state_default_node.name, state_default_node)
+            lg_builder.set_entry_point(state_default_node.name)
+            lg_builder.add_conditional_edges(state_default_node.name, TransitionalEdge(entry_point))
         else:
             # if no state variables are defined, set the entry point directly
             lg_builder.set_entry_point(entry_point)
@@ -619,7 +815,7 @@ def create_graph(
         )
     except ValueError as e:
         raise ValueError(
-            f"Validation of the schema failed. {e}\n\nDEBUG INFO:**Schema Nodes:**\n\n{lg_builder.nodes}\n\n**Schema Enges:**\n\n{lg_builder.edges}\n\n**Tools Available:**\n\n{tools}")
+            f"Validation of the schema failed. {e}\n\nDEBUG INFO:**Schema Nodes:**\n\n*{'\n*'.join(lg_builder.nodes.keys())}\n\n**Schema Edges:**\n\n{lg_builder.edges}\n\n**Tools Available:**\n\n{format_tools(tools)}")
     # If building a nested subgraph, return the raw CompiledStateGraph
     if for_subgraph:
         return graph
@@ -633,6 +829,46 @@ def create_graph(
     )
     return compiled.validate()
 
+def format_tools(tools_list: list) -> str:
+    """Format a list of tool names into a comma-separated string."""
+    try:
+        return ', '.join([tool.name for tool in tools_list])
+    except Exception as e:
+        logger.warning(f"Failed to format tools list: {e}")
+        return str(tools_list)
+
+def set_defaults(d):
+    """Set default values for dictionary entries based on their type."""
+    type_defaults = {
+        'str': '',
+        'list': [],
+        'dict': {},
+        'int': 0,
+        'float': 0.0,
+        'bool': False,
+        # add more types as needed
+    }
+    # Build state_types mapping with STRING type names (not actual type objects)
+    state_types = {}
+
+    for k, v in d.items():
+        # Skip 'input' key as it is not a state initial variable
+        if k == 'input':
+            continue
+        # set value or default if type is defined
+        if 'value' not in v:
+            v['value'] = type_defaults.get(v['type'], None)
+
+        # Also build the state_types mapping with STRING type names
+        var_type = v['type'] if isinstance(v, dict) else v
+        if var_type in ['str', 'int', 'float', 'bool', 'list', 'dict', 'number']:
+            # Store the string type name, not the actual type object
+            state_types[k] = var_type if var_type != 'number' else 'int'
+
+    # Add state_types as a default value that will be set at initialization
+    # Use string type names to avoid serialization issues
+    d['state_types'] = {'type': 'dict', 'value': state_types}
+    return d
 
 def convert_dict_to_message(msg_dict):
     """Convert a dictionary message to a LangChain message object."""
@@ -665,56 +901,208 @@ class LangGraphAgentRunnable(CompiledStateGraph):
     def invoke(self, input: Union[dict[str, Any], Any],
                config: Optional[RunnableConfig] = None,
                *args, **kwargs):
-        logger.info(f"Incomming Input: {input}")
-        if not config.get("configurable", {}).get("thread_id"):
+        logger.info(f"Incoming Input: {input}")
+        if config is None:
+            config = RunnableConfig()
+        if not config.get("configurable", {}).get("thread_id", ""):
             config["configurable"] = {"thread_id": str(uuid4())}
         thread_id = config.get("configurable", {}).get("thread_id")
+        
+        # Check if checkpoint exists early for chat_history handling
+        checkpoint_exists = self.checkpointer and self.checkpointer.get_tuple(config)
+        
         # Handle chat history and current input properly
         if input.get('chat_history') and not input.get('messages'):
-            # 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]
-        
+            if checkpoint_exists:
+                # Checkpoint already has conversation history - discard redundant chat_history
+                input.pop('chat_history', None)
+            else:
+                # No checkpoint - 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
+        # Track if input came from messages to handle content extraction properly
+        input_from_messages = False
+        if not input.get('input'):
+            if input.get('messages'):
+                input['input'] = [next((msg for msg in reversed(input['messages']) if isinstance(msg, HumanMessage)),
+                                      None)]
+                if input['input'] is not None:
+                    input_from_messages = True
+
         # Append current input to existing messages instead of overwriting
         if input.get('input'):
-            current_message = HumanMessage(content=input.get('input'))
+            if isinstance(input['input'], str):
+                current_message = input['input']
+            else:
+                # input can be a list of messages or a single message object
+                current_message = input.get('input')[-1]
+
+            # TODO: add handler after we add 2+ inputs (filterByType, etc.)
+            if isinstance(current_message, HumanMessage):
+                current_content = current_message.content
+                if isinstance(current_content, list):
+                    # Extract text parts and keep non-text parts (images, etc.)
+                    text_contents = []
+                    non_text_parts = []
+                    
+                    for item in current_content:
+                        if isinstance(item, dict) and item.get('type') == 'text':
+                            text_contents.append(item['text'])
+                        elif isinstance(item, str):
+                            text_contents.append(item)
+                        else:
+                            # Keep image_url and other non-text content
+                            non_text_parts.append(item)
+                    
+                    # Set input to the joined text
+                    input['input'] = ". ".join(text_contents) if text_contents else ""
+                    
+                    # If this message came from input['messages'], update or remove it
+                    if input_from_messages:
+                        if non_text_parts:
+                            # Keep the message but only with non-text content (images, etc.)
+                            current_message.content = non_text_parts
+                        else:
+                            # All content was text, remove this message from the list
+                            input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
+                    else:
+                        # Message came from input['input'], not from input['messages']
+                        # If there are non-text parts (images, etc.), preserve them in messages
+                        if non_text_parts:
+                            # Initialize messages if it doesn't exist or is empty
+                            if not input.get('messages'):
+                                input['messages'] = []
+                            # Create a new message with only non-text content
+                            non_text_message = HumanMessage(content=non_text_parts)
+                            input['messages'].append(non_text_message)
+                
+                elif isinstance(current_content, str):
+                    # on regenerate case
+                    input['input'] = current_content
+                    # If from messages and all content is text, remove the message
+                    if input_from_messages:
+                        input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
+                else:
+                    input['input'] = str(current_content)
+                    # If from messages, remove since we extracted the content
+                    if input_from_messages:
+                        input['messages'] = [msg for msg in input['messages'] if msg is not current_message]
+            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:
+                # NOTE: Commented out to prevent duplicates with input['input']
+                # input['messages'] = [current_message]
+        
+        # Validate that input is not empty after all processing
+        if not input.get('input'):
+            raise RuntimeError(
+                "Empty input after processing. Cannot send empty string to LLM. "
+                "This likely means the message contained only non-text content "
+                "with no accompanying text."
+            )
+        
+        logger.info(f"Input: {thread_id} - {input}")
+        try:
+            if self.checkpointer and self.checkpointer.get_tuple(config):
+                if config.pop("should_continue", False):
+                    invoke_input = input
+                else:
+                    self.update_state(config, input)
+                    invoke_input = None
+                result = super().invoke(invoke_input, config=config, *args, **kwargs)
             else:
-                # No existing messages, create new list
-                input['messages'] = [current_message]
-        logging.info(f"Input: {thread_id} - {input}")
-        if self.checkpointer and self.checkpointer.get_tuple(config):
-            self.update_state(config, input)
-            result = super().invoke(None, config=config, *args, **kwargs)
-        else:
-            result = super().invoke(input, config=config, *args, **kwargs)
+                result = super().invoke(input, config=config, *args, **kwargs)
+        except GraphRecursionError as e:
+            current_recursion_limit = config.get("recursion_limit", 0)
+            logger.warning("ToolExecutionLimitReached caught in LangGraphAgentRunnable: %s", e)
+            return self._handle_graph_recursion_error(
+                config=config,
+                thread_id=thread_id,
+                current_recursion_limit=current_recursion_limit,
+            )
+
         try:
-            if self.output_variables and self.output_variables[0] != "messages":
-                # If output_variables are specified, use the value of first one or use the last messages as default
-                output = result.get(self.output_variables[0], result['messages'][-1].content)
+            # Check if printer node output exists
+            printer_output = result.get(PRINTER_NODE_RS)
+            if printer_output == PRINTER_COMPLETED_STATE:
+                # Printer completed, extract last AI message
+                messages = result['messages']
+                output = next(
+                    (msg.content for msg in reversed(messages)
+                     if not isinstance(msg, HumanMessage)),
+                    messages[-1].content
+                ) if messages else result.get('output')
+            elif printer_output is not None:
+                # Printer node has output (interrupted state)
+                output = printer_output
             else:
-                output = result['messages'][-1].content
-        except:
-            output = list(result.values())[-1]
-        thread_id = None
+                # No printer node, extract last AI message from messages
+                messages = result.get('messages', [])
+                output = next(
+                    (msg.content for msg in reversed(messages)
+                     if not isinstance(msg, HumanMessage)),
+                    None
+                )
+        except Exception:
+            # Fallback: try to get last value or last message
+            output = str(list(result.values())[-1]) if result else 'Output is undefined'
         config_state = self.get_state(config)
-        if config_state.next:
-            thread_id = config['configurable']['thread_id']
+        is_execution_finished = not config_state.next
+        if is_execution_finished:
+            thread_id = None
+
+        final_output = f"Assistant run has been completed, but output is None.\nAdding last message if any: {messages[-1] if messages else []}" if is_execution_finished and output is None else output
 
         result_with_state = {
-            "output": output,
+            "output": final_output,
             "thread_id": thread_id,
-            "execution_finished": not config_state.next
+            "execution_finished": is_execution_finished
         }
 
         # Include all state values in the result
         if hasattr(config_state, 'values') and config_state.values:
+            # except of key = 'output' which is already included
+            for key, value in config_state.values.items():
+                if key != 'output':
+                    result_with_state[key] = value
+
+        return result_with_state
+
+    def _handle_graph_recursion_error(
+            self,
+            config: RunnableConfig,
+            thread_id: str,
+            current_recursion_limit: int,
+    ) -> dict:
+        """Handle GraphRecursionError by returning a soft-boundary response."""
+        config_state = self.get_state(config)
+        is_execution_finished = False
+
+        friendly_output = (
+            f"Tool step limit {current_recursion_limit} reached for this run. You can continue by sending another "
+            "message or refining your request."
+        )
+
+        result_with_state: dict[str, Any] = {
+            "output": friendly_output,
+            "thread_id": thread_id,
+            "execution_finished": is_execution_finished,
+            "tool_execution_limit_reached": True,
+        }
+
+        if hasattr(config_state, "values") and config_state.values:
             for key, value in config_state.values.items():
-                result_with_state[key] = value
+                if key != "output":
+                    result_with_state[key] = value
 
         return result_with_state