from llama_index.core import SimpleDirectoryReader
from llama_index.core.extractors import SummaryExtractor,QuestionsAnsweredExtractor
from llama_index.core.node_parser import TokenTextSplitter
from llama_index.core.ingestion import IngestionPipeline, IngestionCache
from llama_index.core.schema import TransformComponent

class CustomTransformation(TransformComponent):
  def __call__(self, nodes, **kwargs):
    # umieść tutaj dowolną logikę przekształacania węzłów
    return nodes
    
reader = SimpleDirectoryReader('pliki')
documents = reader.load_data()
try: 
    cached_hashes = IngestionCache.from_persist_path(
"./ingestion_cache.json"
)
    print("Znaleziono pliki cache. Działanie z użyciem pamięci podręcznej...")
except:
    cached_hashes = ""
    print("Nie znaleziono pliku cache. Działanie bez pamięci podręcznej...")

pipeline = IngestionPipeline(
    transformations = [
        CustomTransformation(),
        TokenTextSplitter(
            separator=" ", 
            chunk_size=512, 
            chunk_overlap=128), 
        SummaryExtractor(), 
        QuestionsAnsweredExtractor(
            questions=3
        )
    ],
    cache=cached_hashes
)

nodes = pipeline.run(documents=documents, show_progress=True)
pipeline.cache.persist("./ingestion_cache.json")

print("Wczytano wszystkie dokumenty")

