LangGraph compatibility and optimization issues

[doxav]

After many tries, I always get some errors when trying to optimize an agent workflow made using langgraph or it does not optimize:

• I could only optimize the main function which generates the graph (e.g. generate_report in the example) but when I create a longer graph it would rarely optimize it.
• I cannot train/bundle function of a graph's node directly (ValueError: no signature found for builtin...) => the workaround I found is to move the code of this function into another function which I set it trainable and call it (e.g. plan_node_train and plan_node in the example) => it does not generate an error but seems not to be optimized.
• only the main function generating the graph (not a node) seems to be optimizable (no node, no function from graph's node is optimized)
• optimizing a trace node value does not seem to work.

You can directly test the code here

from langgraph.graph import StateGraph, START, END
from opto.trace import node, bundle
from opto.optimizers import OptoPrime

state_plan = node("Initial plan: Execute Task A, Task B, and Task C.", trainable=True, description="This represents the current plan of the agent.")

@bundle(trainable=True)
def plan_node_train(state: dict):
    """Creates an initial plan."""
    global state_plan
    state['plan'] = state_plan
    return state

#@bundle(trainable=True)
def plan_node(state: dict):
    return plan_node_train(state)
#    """Creates an initial plan."""
#    state['plan'] = "Initial plan: Execute Task A, Task B, and Task C."
#    return state

@bundle(trainable=True)
def self_critique_node_train(state: dict):
    """Improves the plan based on self-critique."""
    # For illustration, we simply append an improvement comment.
    state['plan'] += " -- Improved after self critique."
    return state

def self_critique_node(state: dict):
    return self_critique_node_train(state)

@bundle(trainable=True)
def finalize_node_train(state: dict):
    """Finalizes the report using the current plan."""
    state['final_report'] = state['plan'] + " -- Final Report."
    return state['final_report']

def finalize_node(state: dict):
    return finalize_node_train(state)

@bundle(trainable=True)
def generate_report():
    # Initialize an empty state
    initial_state = {}

    # Create a simple LangGraph with 3 nodes:
    # START -> plan_node -> self_critique_node -> finalize_node -> END
    graph_builder = StateGraph(dict)
    graph_builder.add_node("plan", plan_node)
    graph_builder.add_node("self_critique", self_critique_node)
    graph_builder.add_node("finalize", finalize_node)

    graph_builder.add_edge(START, "plan")
    graph_builder.add_edge("plan", "self_critique")
    graph_builder.add_edge("self_critique", "finalize")
    graph_builder.add_edge("finalize", END)

    # Compile and invoke the graph
    graph = graph_builder.compile()
    final_state = graph.invoke(initial_state)
    filnal_report_str = f"Final Report: {final_state}"
    print(filnal_report_str)
    return filnal_report_str

parameters =  generate_report.parameters() # [state_plan] + generate_report.parameters() + plan_node_train.parameters() + self_critique_node_train.parameters() + finalize_node_train.parameters()
# parameters = [state_plan] # NO OPTIMIZATION
# parameters = plan_node_train.parameters() # NO OPTIMIZATION
optimizer = OptoPrime(parameters) #, memory_size=2)
optimizer.zero_feedback()

report = generate_report()
# Run a dummy backward pass and step to update parameters.
optimizer.backward(report, "The report quality is low, plan and content are far too short and does not align with research pratice")
optimizer.step()

# Print out optimized parameters (for illustration)
for param in optimizer.parameters:
    print("Optimized parameter:", param.name, param.data)

[allenanie]

Thank you!!! This is a great starting point -- I'll have time to look through it over the weekend

[allenanie]

I went through the example -- this is actually a known issue in Trace's design. We made a choice of not tracing node that is operated inside a function decorated with bundle. What it means is that:

a = node(3)

@bundle(trainalbe=False)
def add1(a):
   """add function"""
   return a + 3      // let's refer to this output as b

def add2(a):
   """add function"""
   return a + 3   // let's refer to this output as b

How would add1 and add2 differ?
They actually create two different graphs!

If you use add1, then the graph is a -> add1 ("add function") -> b
If you use add2, then the graph is a, 3 -> + -> b

bundle hides operations applied in the node and only connects the input with output.

We made this decision to prevent graph blow-up and allow users to summarize operations inside a function.

I suspect this is why generate_report can be updated, but not nodes used inside -- because they are not part of the computational graph.

[allenanie]

We had some mini project with an intern last summer about tracing inside bundle -- but that project might have been put on hold -- we'll figure out a decision on this.

For now, you should write them in a modular fashion.

[doxav]

Thanks for your answer @allenanie .

I understand now what we can't do in a bundle, but I am still not clear how, in the scope/goal of a workflow (generate_report), I can optimize:

• state_plan and/or other prompts nodes : optimizer = OptoPrime([state_plan]) generate an error "TypeError: 'MessageNode' object does not support item assignment" when I comment "@Bundle(trainable=True)" on generate_report and run report = generate_report())
• several functions from the workflow: if I comment "@Bundle(trainable=True)" on generate_report, I get the same error "TypeError: 'MessageNode' object does not support item assignment"

Thanks for your help

[doxav]

Hello,

1. I solved the optimization of multiple agents' prompts in a graph LangGraph (signature using unbind_if_needed, return format of last node). You can look at the solution https://colab.research.google.com/drive/1i7wsnbLSbjuoy2uWNpRf3a7ov7oEYpnf?usp=sharing" => now ir properly optimize different nodes

2. I am still unable to optimize the nodes/agents' functions themselves despite they are not sub functions inside a bundle. I all tried to make the code of those functions a string node which I compile at run time to avoid the bundle hidding mechanism. Please check the full code on the colab

[allenanie]

This is actually super cool and interesting -- can't believe it worked!

I am still enable to optimize the nodes/agents' functions themselves despite they are not sub functions inside a bundle. I all tried to make the code of those functions a string node which I compile at run time to avoid the bundle hidding mechanism. Please check the full code on the colab

I took a quick look at your CoLab, can you point out exactly what section/block I should be looking at? (I assume #3?)

Your

def plan_node(state: dict):
    """Creates an initial plan or improve an existing one if exists"""
    #new_state = state.copy() // new_state['plan'] = state_plan // return new_state
    state['plan'] = state_plan
    return state

This is not wrapped in bundle -- so it shouldn't be updated?

[doxav]

Yes, it's #3. First approach is standard, the other put the code of a function as a string in a node and dynamically compile it. None works, despite the first approach works for prompts optimization.

plan_mode is not wrapped in bundle initially but "dynamically" using this (it is equivalent to decorating and works well, it allows to run different test on the same code):

if not isinstance(plan_node, FunModule): plan_node = bundle(trainable=True)(plan_node)