from typing import List
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
import langchain.docstore.document as docstore
from loguru import logger

from settings import COLLECTION_NAME, PERSIST_DIRECTORY

from .vortex_pdf_parser import VortexPdfParser
from .vortext_content_iterator import VortexContentIterator


class VortexIngester:

    def __init__(self, content_folder: str):
        self.content_folder = content_folder

    def ingest(self) -> None:
        vortex_content_iterator = VortexContentIterator(self.content_folder)
        vortex_pdf_parser = VortexPdfParser()

        chunks: List[docstore.Document] = []
        for document in vortex_content_iterator:
            vortex_pdf_parser.set_pdf_file_path(document)
            document_chunks = vortex_pdf_parser.clean_text_to_docs()
            chunks.extend(document_chunks)
            logger.info(f"Extracted {len(chunks)} chunks from {document}")

        embeddings = OpenAIEmbeddings(client=None)
        logger.info("Loaded embeddings")
        vector_store = Chroma.from_documents(
            chunks,
            embeddings,
            collection_name=COLLECTION_NAME,
            persist_directory=PERSIST_DIRECTORY,
        )

        logger.info("Created Chroma vector store")
        vector_store.persist()
        logger.info("Persisted Chroma vector store")

langchain.embeddings.OpenAIEmbeddings langchain.vectorstores.Chroma.from_documents
[((985, 1014), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'client': 'None'}), '(client=None)\n', (1001, 1014), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1023, 1055), 'loguru.logger.info', 'logger.info', (['"""Loaded embeddings"""'], {}), "('Loaded embeddings')\n", (1034, 1055), False, 'from loguru import logger\n'), ((1079, 1194), 'langchain.vectorstores.Chroma.from_documents', 'Chroma.from_documents', (['chunks', 'embeddings'], {'collection_name': 'COLLECTION_NAME', 'persist_directory': 'PERSIST_DIRECTORY'}), '(chunks, embeddings, collection_name=COLLECTION_NAME,\n    persist_directory=PERSIST_DIRECTORY)\n', (1100, 1194), False, 'from langchain.vectorstores import Chroma\n'), ((1259, 1301), 'loguru.logger.info', 'logger.info', (['"""Created Chroma vector store"""'], {}), "('Created Chroma vector store')\n", (1270, 1301), False, 'from loguru import logger\n'), ((1341, 1385), 'loguru.logger.info', 'logger.info', (['"""Persisted Chroma vector store"""'], {}), "('Persisted Chroma vector store')\n", (1352, 1385), False, 'from loguru import logger\n')]