Add core functionality for LlamaIndex integration

[augray]

Regardless of what integration route we take with LlamaIndex, we will need to be able to upload from their "nodes" abstraction. This PR adds that capability. It turns out some of the auto-inferred types from the object you load from LlamaIndex document nodes like this can't be written to Parquet, so there's some logic to sanitize the data and remove problematic columns. Additionally, the relationships and metadata attributes have some nested data which it would often be useful to have as Airtrain columns, to get insights based on them. This PR contains logic to flatten those attributes into multiple ones.

Testing

Without embeddings, full documents:

Driver script:

import os

from llama_index.readers.github import GithubRepositoryReader, GithubClient

from airtrain import upload_from_llama_nodes


def main() -> None:
    github_token = os.environ.get("GITHUB_TOKEN")
    owner = "sematic-ai"
    repo = "sematic"
    branch = "main"
    github_client = GithubClient(github_token=github_token, verbose=True)
    documents = GithubRepositoryReader(
        github_client=github_client,
        owner=owner,
        repo=repo,
        use_parser=False,
        verbose=False,
        filter_directories=(
            ["docs"],
            GithubRepositoryReader.FilterType.INCLUDE,
        ),
        filter_file_extensions=(
            [
                ".md",
            ],
            GithubRepositoryReader.FilterType.INCLUDE,
        ),
    ).load_data(branch=branch)
    result = upload_from_llama_nodes(
        documents,
        name="Sematic Docs Via Llama Index",
    )
    print(f"Uploaded {result.size} rows to {result.name}. View at: {result.url}")


if __name__ == "__main__":
    main()

Note: implemented before relationship/metadata unpacking was implemented.

resulting dataset

With embeddings, after chunking

import os

from llama_index.readers.github import GithubRepositoryReader, GithubClient
from llama_index.core.node_parser import (
    SentenceSplitter,
    SemanticSplitterNodeParser,
)
from llama_index.embeddings.openai import OpenAIEmbedding

from airtrain import upload_from_llama_nodes


def main() -> None:
    github_token = os.environ.get("GITHUB_TOKEN")
    owner = "sematic-ai"
    repo = "sematic"
    branch = "main"
    github_client = GithubClient(github_token=github_token, verbose=True)
    documents = GithubRepositoryReader(
        github_client=github_client,
        owner=owner,
        repo=repo,
        use_parser=False,
        verbose=False,
        filter_directories=(
            ["docs"],
            GithubRepositoryReader.FilterType.INCLUDE,
        ),
        filter_file_extensions=(
            [
                ".md",
            ],
            GithubRepositoryReader.FilterType.INCLUDE,
        ),
    ).load_data(branch=branch)
    embed_model = OpenAIEmbedding()
    splitter = SemanticSplitterNodeParser(
        buffer_size=1, breakpoint_percentile_threshold=95, embed_model=embed_model
    )
    nodes = splitter.get_nodes_from_documents(documents)
    result = upload_from_llama_nodes(
        nodes,
        name="Sematic Docs Via Llama Index, split + embed",
    )
    print(f"Uploaded {result.size} rows to {result.name}. View at: {result.url}")


if __name__ == "__main__":
    main()

resulting dataset

[augray]

Had to add this because the llama-index package has the <4.0 constraint.

[augray]

Other APIs to initiate this upload may be added in subsequent PR(s), but upload_from_llama_nodes will serve as a foundation for their implementation.