Component#

Component is to LLM task pipelines what nn.Module is to PyTorch models. It is the base class for components such as Prompt, ModelClient, Generator, Retriever in LightRAG. Your task pipeline should also subclass from Component.

Design#

Different from PyTorch’s nn.Module, which works exclusively with Tensor and Parameter to train models with weights and biases, our component can work with different types of data, from a string or a list of strings to a list of Document.

Here is the comparison of writing a PyTorch model and a LightRAG task pipeline.

PyTorch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))
AdalFlow
from adalflow.core import Component, Generator
from adalflow.components.model_client import OpenAIClient


template_doc = r"""<SYS> You are a doctor </SYS> User: {{input_str}}"""

class DocQA(Component):
    def __init__(self):
        super().__init__()
        self.doc = Generator(
            template=template_doc,
            model_client=OpenAIClient(),
            model_kwargs={"model": "gpt-3.5-turbo"},
        )

    def call(self, query: str) -> str:
        return self.doc(prompt_kwargs={"input_str": query}).data

As the fundamental building block in LLM task pipelines, the component is designed to serve five main purposes:

  1. Standardize the interface for all components. This includes the __init__ method, the call method for synchronous calls, the acall method for asynchronous calls, and the __call__ method, which by default calls the call method.

  2. Provide a unified way to visualize the structure of the task pipeline via the __repr__ method. Subclasses can additionally add the _extra_repr method to include more information than the default __repr__ method.

  3. Track and add all subcomponents and parameters automatically and recursively to assist in the building and optimizing process of the task pipeline.

  4. Manage the states and serialization, with state_dict and load_state_dict methods specifically for parameters, and the to_dict method for serialization of all states within the component’s attributes, from subcomponents to parameters, to any other attributes of various data types.

  5. Make all components configurable using `json` or `yaml` files. This is especially useful for experimenting or building data processing pipelines.

These features are key to keeping the LightRAG pipeline transparent, flexible, and easy to use. By subclassing from the Component class, you will get most of these features out of the box.

Component in Action#

In this note, we are creating an AI doctor to answer medical questions. Run the DocQA on a query:

doc = DocQA()
print(doc("What is the best treatment for headache?"))

The response is:

As a doctor, the best treatment for a headache would depend on the underlying cause of the headache. Typically, over-the-counter pain relievers such as acetaminophen, ibuprofen, or aspirin can help to alleviate the pain. However, if the headache is severe or persistent, it is important to see a doctor for further evaluation and to determine the most appropriate treatment option. Other treatment options may include prescription medications, lifestyle modifications, stress management techniques, and relaxation techniques.

Configure from file#

As the above example shows, we added subcomponent via attributes. We can also use methods to add more subcomponnents or parameters.

from adalflow.core.parameter import Parameter

doc.register_parameter("demo", param=Parameter(data="demo"))
# list all parameters
for param in doc.named_parameters():
    print(param)

The output:

('demo', Parameter: demo)

You can easily save the detailed states:

from utils.file_io import save_json

save_json(doc.to_dict(), "doc.json")

To add even more flexibility, we provide FunComponent and Sequential for more advanced use cases.

Serialization and deserialization#

We provide the is_pickable method to check if the component is pickable. It is good practice to ensure that any of your components are pickable.

FunComponent#

Use fun_to_component as a decorator to convert any function to a Component with its unique class name.

FunComponent is a subclass of Component that allows you to define a component with a function. You can directly use this class as:

from adalflow.core.component import FunComponent

def add_one(x):
    return x + 1

fun_component = FunComponent(add_one)
print(fun_component(1))
print(type(fun_component))

The printout:

2
<class 'core.component.FunComponent'>

We also have fun_to_component to convert a function to a FunComponent via a decorator or by directly calling the function. This approach gives you a unique component converted from the function name.

Via direct call:

from adalflow.core.component import fun_to_component

fun_component = fun_to_component(add_one)
print(fun_component(1))
print(type(fun_component))

The output:

2
<class 'adalflow.core.component.AddOneComponent'>

Using a decorator is an even more convenient way to create a component from a function:

@fun_to_component
def add_one(x):
    return x + 1

print(add_one(1))
print(type(add_one))

# output:
# 2
# <class 'adalflow.core.component.AddOneComponent'>

Sequential#

We have the Sequential class, which is similar to PyTorch’s nn.Sequential class. This is especially useful for chaining together components in a sequence, much like the concept of chain or pipeline in other LLM libraries. Let’s put the FunComponent` and DocQA` together in a sequence:

from adalflow.core.container import Sequential

@fun_to_component
def enhance_query(query:str) -> str:
    return query + "Please be concise and only list the top treatments."

seq = Sequential(enhance_query, doc)

query = "What is the best treatment for headache?"
print(seq(query))

We automatically enhance users’ queries before passing them to the DocQA component. The output is:

1. Over-the-counter pain relievers like acetaminophen, ibuprofen, or aspirin
2. Rest and relaxation
3. Stay hydrated and drink plenty of water

The structure of the sequence using print(seq):

Sequential(
(0): EnhanceQueryComponent()
(1): DocQA(
        (doc): Generator(
        model_kwargs={'model': 'gpt-3.5-turbo'}, model_type=ModelType.LLM
        (prompt): Prompt(template: <SYS> You are a doctor </SYS> User: {{input_str}}, prompt_variables: ['input_str'])
        (model_client): OpenAIClient()
        )
    )
)

API reference

We will cover more advanced use cases in the upcoming tutorials.

On this page