lenox-ai/prototype

1from langchain.document_loaders import PyPDFLoader, TextLoader2from langchain.docstore.document import Document3from typing import List4from langchain.text_splitter import (5    RecursiveCharacterTextSplitter,6)7 8 9def load_docs(file_path: str, with_pageinfo: bool = True) -> List[Document]:10    """Load a file and return the text.11 12    Args:13        file_path (str): Path to the pdf file. This can either be a local path or a tempfile.TemporaryFileWrapper_.14        with_pageinfo (bool, optional): If True the page information is added to the document. Defaults to True.15 16    Raises:17        ValueError: If the file type is not supported.18 19    Returns:20        List[Document]: List of documents.21    """22    if file_path.endswith(".pdf"):23        # load documents24        loader = PyPDFLoader(file_path)25        docs = loader.load()26        # # split documents27        # text_splitter = RecursiveCharacterTextSplitter(28        #     chunk_size=1000, chunk_overlap=15029        # )30        # docs = text_splitter.split_documents(docs)31    elif file_path.endswith(".txt"):32        loader = TextLoader(file_path)33        docs = loader.load()34    else:35        raise ValueError(36            f"File type ({file_path.split('.')[1]}) not supported. Please upload a pdf or txt file."37        )38    for doc in docs:39        doc.page_content = doc.page_content.replace("\n", " \n ")40        # if doc contains a page append it to the text41        if with_pageinfo and hasattr(doc, "metadata"):42            doc.page_content = f"(Quelle Seite: {doc.metadata.get('page')+1}) .".join(43                doc.page_content.split(" .")44            )45 46    return docs47