LLMStack/llmstack/data/transformations/__init__.py at 69609aa99c36e75cbe6c4bd5797ec3b3573e271c · trypromptly/LLMStack · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from functools import cache

from llmstack.data.transformations.llamindex.embeddings_generator import (
    EmbeddingsGenerator,
)
from llmstack.data.transformations.llamindex.splitters import (
    CodeSplitter,
    SemanticDoubleMergingSplitterNodeParser,
    SemanticSplitterNodeParser,
    SentenceSplitter,
    SentenceWindowNodeParser,
    TokenTextSplitter,
)
from llmstack.data.transformations.splitters import CSVTextSplitter
from llmstack.data.transformations.unstructured.splitters import UnstructuredIOSplitter


@cache
def get_transformer_cls(slug, provider_slug):
    for cls in [UnstructuredIOSplitter, EmbeddingsGenerator, CSVTextSplitter]:
        if cls.slug() == slug and cls.provider_slug() == provider_slug:
            return cls

    return None


__all__ = [
    "CodeSplitter",
    "SemanticDoubleMergingSplitterNodeParser",
    "SemanticSplitterNodeParser",
    "SentenceSplitter",
    "SentenceWindowNodeParser",
    "TokenTextSplitter",
    "CSVTextSplitter",
    "UnstructuredIOSplitter",
    "EmbeddingsGenerator",
]