Enhancing LlamaIndex

Introduction to Custom Node Types in LlamaIndex

LlamaIndex is a powerful framework for building LLM-powered applications, and one of its key features is the ability to create custom node types. These custom nodes allow developers to extend the functionality of LlamaIndex and tailor it to specific use cases. In this blog post, we'll explore how to implement custom node types in Python using LlamaIndex.

Why Use Custom Node Types?

Before we dive into the implementation, let's understand why custom node types are valuable:

1. Flexibility: Custom nodes allow you to represent data in ways that best suit your application's needs.
2. Enhanced functionality: You can add specific methods and attributes to your nodes, improving data processing capabilities.
3. Better organization: Custom nodes help in structuring complex data hierarchies more effectively.

Creating a Custom Node Type

To create a custom node type, we need to subclass the IndexNode or TextNode class from LlamaIndex. Here's a basic example:

from llama_index.data_structs.node import Node

class CustomNode(Node):
    def __init__(self, text, extra_info=None, **kwargs):
        super().__init__(text, **kwargs)
        self.extra_info = extra_info

    def get_extra_info(self):
        return self.extra_info

In this example, we've created a CustomNode that includes an extra_info attribute. This allows us to store additional information alongside the node's text content.

Implementing Custom Methods

One of the advantages of custom nodes is the ability to add specific methods. Let's enhance our CustomNode with a method to process its content:

import nltk
from llama_index.data_structs.node import Node

class CustomNode(Node):
    def __init__(self, text, extra_info=None, **kwargs):
        super().__init__(text, **kwargs)
        self.extra_info = extra_info

    def get_extra_info(self):
        return self.extra_info

    def get_key_phrases(self, num_phrases=3):
        sentences = nltk.sent_tokenize(self.text)
        words = nltk.word_tokenize(self.text)
        tagged = nltk.pos_tag(words)
        
        noun_phrases = []
        for word, pos in tagged:
            if pos.startswith('NN'):
                noun_phrases.append(word)
        
        return noun_phrases[:num_phrases]

This get_key_phrases method uses NLTK to extract key noun phrases from the node's text content, which can be useful for summarization or topic extraction tasks.

Using Custom Nodes in LlamaIndex

Now that we have our custom node type, let's see how to use it within LlamaIndex:

from llama_index import GPTSimpleVectorIndex, SimpleDirectoryReader
from llama_index.node_parser import SimpleNodeParser

# Create a custom node parser
class CustomNodeParser(SimpleNodeParser):
    def get_nodes_from_documents(self, documents):
        nodes = super().get_nodes_from_documents(documents)
        return [CustomNode(node.text, extra_info="Custom data") for node in nodes]

# Load documents
documents = SimpleDirectoryReader('data').load_data()

# Parse documents into custom nodes
parser = CustomNodeParser()
nodes = parser.get_nodes_from_documents(documents)

# Create index with custom nodes
index = GPTSimpleVectorIndex(nodes)

# Query the index
response = index.query("What are the main topics?")
print(response)

# Use custom node method
for node in nodes:
    print(f"Key phrases: {node.get_key_phrases()}")
    print(f"Extra info: {node.get_extra_info()}")

In this example, we've created a custom node parser that converts standard nodes into our CustomNode type. We then use these custom nodes to create an index and perform queries.

Advanced Custom Node Types

As you become more comfortable with custom nodes, you can create more complex types to handle specific data structures or processing requirements. For instance, you might create a TimeseriesNode for time-series data or an ImageNode for image-based information.

Here's a quick example of a TimeseriesNode:

import pandas as pd
from llama_index.data_structs.node import Node

class TimeseriesNode(Node):
    def __init__(self, text, timestamp, value, **kwargs):
        super().__init__(text, **kwargs)
        self.timestamp = pd.to_datetime(timestamp)
        self.value = value

    def to_series(self):
        return pd.Series({
            'timestamp': self.timestamp,
            'value': self.value,
            'text': self.text
        })

    def resample(self, rule='D'):
        series = self.to_series()
        return series.resample(rule).mean()

This TimeseriesNode allows you to work with time-based data more effectively within LlamaIndex, including resampling and conversion to pandas Series.

Conclusion

Custom node types in LlamaIndex offer a powerful way to extend the framework's capabilities and tailor it to your specific needs. By implementing custom nodes, you can enhance data processing, improve organization, and create more sophisticated LLM applications. As you continue to work with LlamaIndex, experiment with different custom node types to unlock the full potential of your data and LLM models.