Annotate a class method with @task (Functional API)

[wjaskowski]

Is there a way to annotate a method with @task or @workflow? I'd like to encapsulate the whole workflow in a class for better modularity of my code base but could not find a way to annotate methods. Example code:

from langgraph.func import entrypoint as workflow, task

class MockLLM:
    def generate(self, prompt: str) -> str:
        return f"[MOCKED] Response for: {prompt}"

class Workflow:
    def __init__(self, param: int = 0):
        self.llm = MockLLM()
        self.param = param

    @task
    def call_llm(self, prompt: str) -> str:
        return self.llm.generate(prompt)

    @workflow()
    def simple_workflow(self, topic: str) -> dict:
        if self.param == 0:
            prompt = f"Explain the significance of {topic} in one sentence."
        else:
            prompt = f"Don't do anything"
        result = self.call_llm(prompt).result()
        return {"topic": topic, "explanation": result}

    def invoke(self, topic: str) -> dict:
        return self.simple_workflow.invoke(topic)

if __name__ == "__main__":
    workflow = Workflow()

    output = workflow.invoke("LangGraph")
    print(output)

doesn't work because of

.venv/lib/python3.11/site-packages/langgraph/utils/runnable.py", line 365, in invoke
    ret = context.run(self.func, *args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: Workflow.simple_workflow() missing 1 required positional argument: 'topic'
During task with name 'simple_workflow' and id 'f0f2b29d-9060-5e3d-ec01-cc6aca14df08'

[dionysis11]

Hello! I did a quick research and I found that you can't directly decorate class methods with @task or @workflow decorators because they can't properly handle the self parameter. When Python applies these decorators at class definition time, they don't have context about class instances.

Instead, create your tasks and workflows inside the constructor like this:

def __init__(self):
    # The key insight: create tasks by passing methods to the decorator functions
    self.call_llm_task = task(self.call_llm)
    self.workflow_entry = entrypoint()(self.simple_workflow)

This way, the instance methods are already bound to self when you pass them to the task/workflow functions. It's a bit more code but ensures everything works properly with your class architecture.

Here's a full working example:

from langgraph.func import entrypoint, task

class MockLLM:
    def generate(self, prompt: str) -> str:
        return f"[MOCKED] Response for: {prompt}"

class Workflow:
    def __init__(self, param: int = 0):
        self.llm = MockLLM()
        self.param = param
        # Create tasks by passing methods directly to task() function
        self.call_llm_task = task(self.call_llm, name="call_llm")
        # Create entrypoint by passing method directly
        self.simple_workflow_entry = entrypoint()(self.simple_workflow)

    def call_llm(self, prompt: str) -> str:
        return self.llm.generate(prompt)

    def simple_workflow(self, topic: str) -> dict:
        if self.param == 0:
            prompt = f"Explain the significance of {topic} in one sentence."
        else:
            prompt = f"Don't do anything"
        result = self.call_llm_task(prompt).result()
        return {"topic": topic, "explanation": result}

    def invoke(self, topic: str) -> dict:
        return self.simple_workflow_entry.invoke(topic)

if __name__ == "__main__":
    workflow = Workflow()
    output = workflow.invoke("LangGraph")
    print(output) 

Please correct me if I'm wrong!

[wjaskowski]

I see. Thank you for the workaround! But this is not a limitation of Python, is it? It would be possible to extend the implementation of @task and @workflow to handle also methods, right?

[hinthornw]

Thanks for the request! We'll try to add support

[wjaskowski]

Nice! Is there any place I could track progress for this feature?

[dionysis11]

Hi @hinthornw , I wanted to check if there’s been any progress. If so, I’d be happy to assist. If not, would it be possible for me to be assigned to this?