import inspect
from pathlib import Path
from typing import List

from langchain.chains import LLMChain
from langchain.chat_models.base import BaseChatModel
from langchain.prompts import PromptTemplate


def get_documents(file_path: Path, llm: BaseChatModel):
    file_extension = file_path.suffix
    loader_class_name = get_best_loader(file_extension, llm)
    print(f"loader selected {loader_class_name} for {file_path}")

    if loader_class_name == "None":
        raise Exception(f"No loader found for {file_extension} files.")

    loader_class = get_loader_class(loader_class_name)
    loader = loader_class(str(file_path))
    return loader.load()


def get_loader_class(loader_class_name: str):
    import langchain.document_loaders

    loader_class = getattr(langchain.document_loaders, loader_class_name)
    return loader_class


def get_best_loader(file_extension: str, llm: BaseChatModel):
    loaders = get_loaders()
    prompt = PromptTemplate(
        input_variables=["file_extension", "loaders"],
        template="""
    Among the following loaders, which is the best to load a "{file_extension}" file? \
        Only give me one the class name without any other special characters. If no relevant loader is found, respond "None".

    Loaders: {loaders}
    """,
    )
    chain = LLMChain(llm=llm, prompt=prompt, output_key="loader_class_name")

    return chain({"file_extension": file_extension, "loaders": loaders})["loader_class_name"]


def get_loaders() -> List[str]:
    import langchain_community.document_loaders

    loaders = []
    for _, obj in inspect.getmembers(langchain_community.document_loaders):
        if inspect.isclass(obj):
            loaders.append(obj.__name__)
    return loaders

langchain.chains.LLMChain langchain.prompts.PromptTemplate
[((946, 1275), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['file_extension', 'loaders']", 'template': '"""\n    Among the following loaders, which is the best to load a "{file_extension}" file?         Only give me one the class name without any other special characters. If no relevant loader is found, respond "None".\n\n    Loaders: {loaders}\n    """'}), '(input_variables=[\'file_extension\', \'loaders\'], template=\n    """\n    Among the following loaders, which is the best to load a "{file_extension}" file?         Only give me one the class name without any other special characters. If no relevant loader is found, respond "None".\n\n    Loaders: {loaders}\n    """\n    )\n', (960, 1275), False, 'from langchain.prompts import PromptTemplate\n'), ((1303, 1367), 'langchain.chains.LLMChain', 'LLMChain', ([], {'llm': 'llm', 'prompt': 'prompt', 'output_key': '"""loader_class_name"""'}), "(llm=llm, prompt=prompt, output_key='loader_class_name')\n", (1311, 1367), False, 'from langchain.chains import LLMChain\n'), ((1581, 1637), 'inspect.getmembers', 'inspect.getmembers', (['langchain_community.document_loaders'], {}), '(langchain_community.document_loaders)\n', (1599, 1637), False, 'import inspect\n'), ((1650, 1670), 'inspect.isclass', 'inspect.isclass', (['obj'], {}), '(obj)\n', (1665, 1670), False, 'import inspect\n')]