from PyPDF2 import PdfReader

import os

import pandas as pd

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.chains.question_answering import load_qa_chain
from langchain.llms import OpenAI
from langchain.cache import InMemoryCache
import langchain

from ocr.ocr import process_pdf
from docs.search import doc_search
from extract.extract import knowledge_graph
from data.processing import rowify
from edi.edi_formatter import pandas_to_edi

from openai.error import InvalidRequestError


langchain.llm_cache = InMemoryCache()
llm = OpenAI(temperature=0)

embeddings = OpenAIEmbeddings()

chain = load_qa_chain(llm, chain_type="stuff")

pdf_inputs = []

key_ids = ""

query = f"""
Using the unique count of {key_ids} in this document, do the following:
    For each {key_ids}, extract the following information corresponding to the {key_ids}:
    """

rules_template = f"""Just give me the answer with {key_ids} line separated and nothing else."""

pdf_data = []

pdf_dir = '/pdfs/'
fils = os.listdir(pdf_dir)

for fil in fils:
    print("processing: " + fil)
    try:
        pdf_file = pdf_dir + fil
        texts = process_pdf(pdf_file)
        docsearch = doc_search(texts, embeddings)
        hwb_data = knowledge_graph(
            key_id="",
            docsearch=docsearch,
            pdf_inputs=pdf_inputs,
            query=query,
            rules_template=rules_template,
            chain=chain
        )

        mwb = fil.split('-')[1]
        rows = rowify(hwb_data, extra=[mwb])

        pdf_data.extend(rows)
    except InvalidRequestError:
        print(fil, "File needs handler.")

cols = []
df = pd.DataFrame(columns=cols, data=pdf_data)

edi_data = pandas_to_edi(
    edi_type='211',
    df=df,
    edi_key_col="",
    edi_data_col="",
)
langchain.chains.question_answering.load_qa_chain langchain.cache.InMemoryCache langchain.llms.OpenAI langchain.embeddings.openai.OpenAIEmbeddings
[((634, 649), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (647, 649), False, 'from langchain.cache import InMemoryCache\n'), ((656, 677), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)'}), '(temperature=0)\n', (662, 677), False, 'from langchain.llms import OpenAI\n'), ((692, 710), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (708, 710), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((720, 758), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['llm'], {'chain_type': '"""stuff"""'}), "(llm, chain_type='stuff')\n", (733, 758), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((1113, 1132), 'os.listdir', 'os.listdir', (['pdf_dir'], {}), '(pdf_dir)\n', (1123, 1132), False, 'import os\n'), ((1741, 1782), 'pandas.DataFrame', 'pd.DataFrame', ([], {'columns': 'cols', 'data': 'pdf_data'}), '(columns=cols, data=pdf_data)\n', (1753, 1782), True, 'import pandas as pd\n'), ((1795, 1864), 'edi.edi_formatter.pandas_to_edi', 'pandas_to_edi', ([], {'edi_type': '"""211"""', 'df': 'df', 'edi_key_col': '""""""', 'edi_data_col': '""""""'}), "(edi_type='211', df=df, edi_key_col='', edi_data_col='')\n", (1808, 1864), False, 'from edi.edi_formatter import pandas_to_edi\n'), ((1241, 1262), 'ocr.ocr.process_pdf', 'process_pdf', (['pdf_file'], {}), '(pdf_file)\n', (1252, 1262), False, 'from ocr.ocr import process_pdf\n'), ((1283, 1312), 'docs.search.doc_search', 'doc_search', (['texts', 'embeddings'], {}), '(texts, embeddings)\n', (1293, 1312), False, 'from docs.search import doc_search\n'), ((1332, 1463), 'extract.extract.knowledge_graph', 'knowledge_graph', ([], {'key_id': '""""""', 'docsearch': 'docsearch', 'pdf_inputs': 'pdf_inputs', 'query': 'query', 'rules_template': 'rules_template', 'chain': 'chain'}), "(key_id='', docsearch=docsearch, pdf_inputs=pdf_inputs,\n    query=query, rules_template=rules_template, chain=chain)\n", (1347, 1463), False, 'from extract.extract import knowledge_graph\n'), ((1590, 1619), 'data.processing.rowify', 'rowify', (['hwb_data'], {'extra': '[mwb]'}), '(hwb_data, extra=[mwb])\n', (1596, 1619), False, 'from data.processing import rowify\n')]