A chatbot for your documents
#!/usr/bin/env python3

"""
This module provides the entry point for the DocuBot application. It includes functions for building a knowledge base, initializing the application, and running the chatbot. The main function `docubot()` creates an argparse to parse the data directory as a mandatory argument, loads API keys, creates or loads an existing index, and instantiates the QA chatbot. 
"""

from typing import List, TypeVar

T = TypeVar("T")


def chatbot(vector_store):
    """
    A chatbot function that allows users to ask questions and receive answers with citations.

    Args:
        vector_store (list): A list of vectors used for semantic search.

    Returns:
        None
    """
    import time
    from docubot_utils.docubot_utils import answer_question_session, format_citations

    i = 1
    print("Welcome to the chatbot. Type 'quit' or 'exit' to exit")
    chat_history = []
    while True:
        question = input(f"Question #{i}: ")
        if question.lower() in ["quit", "exit"]:
            print("Exiting chatbot")
            time.sleep(2)
            break
        result, chat_history = answer_question_session(
            question, vector_store=vector_store, chat_history=chat_history
        )
        print(f"\nAnswer #{i}: {result['answer']}")
        # Join vector in result['source_document'] into a string seperated by newline character
        citations = format_citations(result["source_documents"])
        print(f"\nCitations #{i}:\n{citations}")
        print(f"\n {'-' * 50} \n")
        i += 1


def build_kb(data_directory: str) -> List[T]:
    """
    Builds a knowledge base by loading documents from the given data directory.

    Args:
        data_directory (str): The path to the directory containing the documents.

    Returns:
        List[T]: A list of chunks, where each chunk contains 512 tokens.
    """
    import os
    from tqdm import tqdm
    from document_loaders.document_loaders import load_document, chunk_data

    list_of_doc_names = []
    # List of supported document types
    extensions = (".pdf", ".docx", ".md", ".txt")
    for root, _, files in os.walk(data_directory):
        for file in files:
            if file.endswith(extensions):
                list_of_doc_names.append(os.path.join(root, file))
            else:
                continue

    # This is the knowledge base that is chunked into 512 tokens per chunk
    data = []
    print(f"Splitting documents into chunks")
    for doc_name in tqdm(list_of_doc_names):
        data.extend(load_document(doc_name))

    print(f"There are {len(data)} pages in the knowledge base")
    chunks = chunk_data(data, chunk_size=512, chunk_overlap=20)
    print(f"These have been split into {len(chunks)} chunks for indexing")
    return chunks


def init():
    """
    Initializes Pinecone API with API key and environment variables.
    """
    from dotenv import load_dotenv, find_dotenv
    import pinecone
    import os

    # Load API keys
    load_dotenv(find_dotenv(), override=True)
    pinecone.init(
        api_key=os.environ.get("PINECONE_API_KEY"),
        environment=os.environ.get("PINECONE_ENV"),
    )


def docubot():
    """
    This function is the entry point of the DocuBot application. It creates an argparse to parse data directory which is a mandatory argument.
    It loads API keys and creates or loads existing index and instantiate QA chatbot.
    """
    import argparse
    import pinecone
    from pinecone_utils.pinecone_utils import (
        create_index,
        create_vector_store,
        fetch_vector_store,
        delete_pinecone_index,
    )
    from langchain.vectorstores import Pinecone

    parser = argparse.ArgumentParser(description="DocuBot")
    parser.add_argument("docs_dir", type=str, help="Path to the documents directory")
    parser.add_argument(
        "--index_name",
        type=str,
        default="docubot-index",
        help="Descriptive name for backend vector db index",
    )
    parser.add_argument(
        "--reload",
        type=str,
        help="Name of index to reload. If 'all', all indices will be deleted.",
    )
    args = parser.parse_args()
    docs_directory = args.docs_dir
    index_name = args.index_name
    reload_index_name = args.reload
    print(f"Instantiating DocuBot for {docs_directory}")
    init()

    vector_store: Pinecone = None
    if reload_index_name:
        # Delete all pinecone indices. As a free user, we can only have one index at a time.
        delete_pinecone_index(reload_index_name)
        create_index(index_name)
        vector_store = create_vector_store(index_name, chunks=build_kb(docs_directory))
    elif index_name in pinecone.list_indexes():
        vector_store = fetch_vector_store(index_name)
    else:
        create_index(index_name)
        vector_store = create_vector_store(index_name, chunks=build_kb(docs_directory))
    chatbot(vector_store=vector_store)


if __name__ == "__main__":
    docubot()

docubot.py
bshastry/docubot
#!/usr/bin/env python3

"""
This module provides functions for answering questions using retrieval-based question answering systems.

It includes functions for answering a single question, answering a question within a conversational context,
and formatting source document citations.

Note: This module requires the 'langchain' library to be installed.
"""
from typing import List, TypeVar

T = TypeVar("T")


def answer_question(question: str, vector_store, num_neighbors=5):
    """
    This function takes a question as input and returns an answer using a retrieval-based question answering system.

    Args:
    - question (str): The question to be answered.
    - vector_store: A vector store object used for retrieving similar documents.
    - num_neighbors (int): The number of similar documents to retrieve.

    Returns:
    - The answer to the input question.
    """
    from langchain.chat_models import ChatOpenAI
    from langchain.chains import RetrievalQA

    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
    retriever = vector_store.as_retriever(
        search_type="similarity", search_kwargs={"k": num_neighbors}
    )
    chain = RetrievalQA.from_chain_type(
        llm=llm, chain_type="stuff", retriever=retriever
    )
    return chain.run(question)


def answer_question_session(
    question: str, vector_store, num_neighbors=5, chat_history=[]
):
    """
    This function takes a question, a vector store, and an optional number of neighbors and chat history.
    It uses the LangChain library to create a conversational retrieval chain, which is used to generate an answer to the question.
    The answer is then appended to the chat history.
    The function returns a dictionary containing the answer and the source documents used to generate the answer, as well as the updated chat history.

    Args:
        question (str): The question to be answered.
        vector_store: The vector store used to retrieve similar documents.
        num_neighbors (int, optional): The number of similar documents to retrieve. Defaults to 5.
        chat_history (list, optional): A list of tuples containing the question and answer pairs from previous conversations. Defaults to an empty list.

    Returns:
        tuple: A tuple containing a dictionary with the answer and source documents, and the updated chat history.
    """
    from langchain.chat_models import ChatOpenAI
    from langchain.chains import ConversationalRetrievalChain

    llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0.7)
    retriever = vector_store.as_retriever(
        search_type="similarity", search_kwargs={"k": num_neighbors}
    )
    conversational_chain = ConversationalRetrievalChain.from_llm(
        llm=llm, retriever=retriever, return_source_documents=True
    )
    result = conversational_chain({"question": question, "chat_history": chat_history})
    chat_history.append((question, result["answer"]))
    return result, chat_history


def format_citations(source_docs: List[T]) -> str:
    """
    Formats a list of source documents into a string of citations.

    Args:
        source_docs (List[T]): A list of source documents.

    Returns:
        str: A string of formatted citations.
    """
    # Iterate through source_docs and reference a field in each index called metadata
    citation_list = []
    for index, doc in enumerate(source_docs):
        source_document_name = doc.metadata["source"]
        page = doc.metadata.get("page")
        page_str = str(int(page)) if page else "not available"
        citation_list.append(
            "{}: {}, page {}".format(index, source_document_name, page_str)
        )
    return "\n".join(citation_list)

docubot_utils/docubot_utils.py
bshastry/docubot
#!/usr/bin/env python3

"""
This module provides functions for loading various types of documents.

It includes functions for loading PDF documents, .docx documents, markdown documents,
plain text documents, documents from Wikipedia, and documents from files or URLs.

Note: This module requires the 'langchain' library to be installed.
"""

from typing import List, TypeVar

T = TypeVar("T")


def load_pdf_document(document_name: str, verbose: bool = False) -> List[T]:
    """
    Load a PDF document using PyPDFLoader.

    Args:
        document_name (str): The name of the PDF document to load.
        verbose (bool, optional): Whether to print a message indicating that the document is being loaded. Defaults to False.

    Returns:
        List[T]: A list of objects representing the loaded PDF document.
    """
    from langchain.document_loaders import PyPDFLoader

    if verbose:
        print(f"Loading {document_name}")
    return PyPDFLoader(document_name).load()


def load_docx_document(document_name: str, verbose: bool = False) -> List[T]:
    """
    Load a document in .docx format using the Docx2txtLoader class.

    Args:
        document_name (str): The name of the document to load.
        verbose (bool, optional): Whether to print verbose output. Defaults to False.

    Returns:
        List[T]: A list of strings representing the document's contents.
    """
    from langchain.document_loaders import Docx2txtLoader

    if verbose:
        print(f"Loading {document_name}")
    return Docx2txtLoader(document_name).load()


def load_markdown_document(document_name: str, verbose: bool = False) -> List[T]:
    """
    Load a markdown document using the UnstructuredMarkdownLoader.

    Args:
        document_name (str): The name of the markdown document to load.
        verbose (bool, optional): Whether to print verbose output. Defaults to False.

    Returns:
        List[T]: A list of objects loaded from the markdown document.
    """
    from langchain.document_loaders import UnstructuredMarkdownLoader

    if verbose:
        print(f"Loading {document_name}")
    return UnstructuredMarkdownLoader(document_name).load()


def load_txt_document(document_name: str, verbose: bool = False) -> List[T]:
    """
    Load a text document using the TextLoader class from langchain.document_loaders.

    Args:
        document_name (str): The name of the document to load.
        verbose (bool, optional): Whether to print a message indicating that the document is being loaded. Defaults to False.

    Returns:
        List[T]: A list of the loaded document's contents.
    """
    from langchain.document_loaders import TextLoader

    if verbose:
        print(f"Loading {document_name}")
    return TextLoader(document_name).load()


def load_from_wikipedia(
    query: str, lang: str = "en", load_max_docs: int = 2
) -> List[T]:
    """
    Load documents from Wikipedia based on the given query and language.

    Args:
        query (str): The query to search for on Wikipedia.
        lang (str, optional): The language of the Wikipedia to search. Defaults to "en".
        load_max_docs (int, optional): The maximum number of documents to load. Defaults to 2.

    Returns:
        List[T]: A list of documents loaded from Wikipedia.
    """
    from langchain.document_loaders import WikipediaLoader

    return WikipediaLoader(query=query, lang=lang, load_max_docs=load_max_docs).load()


def load_document(document_name: str) -> List[T]:
    """
    Load a document from a file or URL.

    Args:
        document_name (str): The name of the document file or URL.

    Returns:
        List[T]: A list of the loaded document's contents.

    Raises:
        None.
    """
    import os
    import validators
    from text_utils.text_utils import return_url_extension

    if validators.url(document_name):
        if return_url_extension(document_name) == ".pdf":
            return load_pdf_document(document_name)
        else:
            print("Unsupported file extension")
            return None
    else:
        _, extension = os.path.splitext(document_name)
        if extension == ".pdf":
            return load_pdf_document(document_name)
        elif extension == ".docx":
            return load_docx_document(document_name)
        elif extension == ".md":
            return load_markdown_document(document_name)
        elif extension == ".txt":
            return load_txt_document(document_name)
        else:
            print("Unsupported file extension")
            return None


def chunk_data(
    data: List[T], chunk_size: int = 512, chunk_overlap: int = 20
) -> List[T]:
    """
    Splits the input data into chunks of specified size using a RecursiveCharacterTextSplitter.

    Args:
        data (str): The input data to be split into chunks.
        chunk_size (int, optional): The size of each chunk. Defaults to 512.

    Returns:
        List[T]: A list of chunks, where each chunk is of type T.
    """
    from text_utils.text_utils import tiktoken_len
    from langchain.text_splitter import RecursiveCharacterTextSplitter

    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=tiktoken_len,
        separators=["\n\n", "\n", " ", ""],
    )
    chunks = text_splitter.split_documents(data)
    return chunks


def merge_document(document: List[T]) -> str:
    """
    Merge a list of documents into a single string.

    Args:
        document (List[T]): A list of documents to merge.

    Returns:
        str: A single string containing the merged documents.
    """
    return "\n\n".join([page.page_content for page in document])

document_loaders/document_loaders.py
bshastry/docubot
#!/usr/bin/env python3
"""
This script provides functionality for summarizing a given document using OpenAI's GPT-3.5-turbo model.
It includes the 'init()' function to initialize environment variables, the 'summarize()' function to generate summaries,
and the 'docusense()' function as the entry point for the script. The 'docusense()' function takes command-line arguments
for the document path, chunk size, and chunk overlap. It utilizes prompts and chains to perform the summarization process.
"""


def init():
    """
    Initializes the environment variables by loading the .env file.

    Returns:
    None
    """
    from dotenv import load_dotenv, find_dotenv

    load_dotenv(find_dotenv(), override=True)


def summarize(
    document: str,
    summary_file: str,
    chunk_size: int,
    chunk_overlap: int,
    max_single_shot_num_tokens: int = 2048,
) -> None:
    """
    Summarizes a given document using OpenAI's GPT-3.5-turbo model.

    Args:
        document (str): The path to the document to be summarized.
        chunk_size (int): The size of each chunk of the document to be summarized.
        chunk_overlap (int): The amount of overlap between each chunk of the document.
        max_single_shot_num_tokens (int, optional): The maximum number of tokens allowed for a single-shot summarization. Defaults to 2048.

    Returns:
        None

    Raises:
        FileNotFoundError: If the specified document path does not exist.
    """
    from langchain.chat_models import ChatOpenAI
    from langchain import PromptTemplate
    from langchain.chains import LLMChain
    from langchain.chains.summarize import load_summarize_chain
    from document_loaders.document_loaders import (
        load_document,
        merge_document,
        chunk_data,
    )
    from text_utils.text_utils import num_tokens_and_cost

    llm = ChatOpenAI(temperature=0, model="gpt-3.5-turbo")

    map_prompt = """
Write a concise summary of the following:
Text: `{text}`
CONCISE SUMMARY:
"""
    map_prompt_template = PromptTemplate(template=map_prompt, input_variables=["text"])
    combine_prompt = """
Write a concise summary of the following text that covers key points.
Add a title to the summary.
Start the summary with an INTRODUCTION PARAGRAPH that gives an overview of the topic FOLLOWED
by BULLET POINTS if possible AND end the summary with a CONCLUSION.
Text: `{text}`
"""
    combine_prompt_template = PromptTemplate(
        template=combine_prompt, input_variables=["text"]
    )

    doc = load_document(document)
    num_tokens, cost = num_tokens_and_cost(doc)
    print(f"Approximate summarization cost: ${cost:.4f}")
    if num_tokens <= max_single_shot_num_tokens:
        chain = LLMChain(llm=llm, prompt=combine_prompt_template)
        print("Running single-shot summarization")
        summary = chain.run({"text": merge_document(doc)})
        print(f"Writing summary to {summary_file}... ", end="")
        with open(summary_file, "w") as f:
            f.write(summary)
        print("Done")
    else:
        chain = load_summarize_chain(
            llm=llm,
            chain_type="map_reduce",
            map_prompt=map_prompt_template,
            combine_prompt=combine_prompt_template,
        )
        print("Running multi-shot summarization")
        chain = load_summarize_chain(
            llm=llm,
            chain_type="map_reduce",
            map_prompt=map_prompt_template,
            combine_prompt=combine_prompt_template,
        )
        summary = chain.run(
            chunk_data(data=doc, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        )
        print(f"Writing summary to {summary_file}... ", end="")
        with open(summary_file, "w") as f:
            f.write(summary)
        print("Done")


def docusense() -> None:
    """
    This function takes in a document path and summary file path and summarizes it using DocuSense.
    It also takes in optional arguments for chunk size and overlap.

    Returns:
        None
    """
    import argparse

    parser = argparse.ArgumentParser(description="DocuSense")
    parser.add_argument(
        "document", type=str, help="Path to the document to be summarized."
    )
    parser.add_argument(
        "summary_file",
        type=str,
        help="Path to the file where summary will be written to.",
    )
    parser.add_argument(
        "--chunk_size", type=int, default=3300, help="Chunk size in tokens."
    )
    parser.add_argument(
        "--chunk_overlap", type=int, default=100, help="Chunk overlap in tokens."
    )
    args = parser.parse_args()
    document = args.document
    summary_file = args.summary_file
    chunk_size = args.chunk_size
    chunk_overlap = args.chunk_overlap
    print(f"Instantiating DocuSense for {document}")
    init()
    try:
        summarize(document, summary_file, chunk_size, chunk_overlap)
    except FileNotFoundError:
        print(f"File {document} not found")


if __name__ == "__main__":
    docusense()

docusense.py
bshastry/docubot
#!/usr/bin/env python3

"""
This module provides functions for interacting with Pinecone indexes.

It includes functions for creating an index, populating it with embeddings,
fetching embeddings from an existing index, and deleting an index.

Note: This module requires the 'pinecone' and 'langchain' libraries to be installed.
"""

from typing import List, TypeVar
from langchain.vectorstores import Pinecone

T = TypeVar("T")


def create_index(index_name: str) -> None:
    """
    Creates a Pinecone index with the given name.

    Args:
        index_name (str): The name of the index to create.

    Raises:
        ValueError: If the index name is already taken.
        pinecone.exceptions.PineconeException: If there is an error creating the index.
    """
    import pinecone

    if index_name in pinecone.list_indexes():
        raise ValueError(f"Index {index_name} already exists.")
    try:
        print(f"Creating index {index_name}... ", end="")
        pinecone.create_index(name=index_name, metric="cosine", dimension=1536)
        print(f"Done")
    except pinecone.exceptions.PineconeException as e:
        raise pinecone.exceptions.PineconeException(f"Error creating index: {str(e)}")


def create_vector_store(index_name: str, chunks: List[T]) -> Pinecone:
    """
    Populates a Pinecone index with embeddings for the input documents (chunks).

    Args:
        index_name (str): The name of the index to populate.
        chunks (List[T]): A list of data chunks to index.

    Returns:
        Pinecone: A Pinecone vector store object containing the indexed data.

    Raises:
        ValueError: If the index name does not exist.
        pinecone.exceptions.PineconeException: If there is an error indexing the documents.
    """
    import pinecone
    from langchain.vectorstores import Pinecone
    from langchain.embeddings.openai import OpenAIEmbeddings
    from text_utils.text_utils import num_tokens_and_cost

    num_tokens, cost = num_tokens_and_cost(chunks)
    # Prompt user whether they want to continue, quit if they don't
    while True:
        user_input = input(
            f"Cost Estimate: ${cost:.4f} for {num_tokens} tokens\n"
            f"Would you like to continue? (y/n)\n"
        )
        if user_input.lower() == "y":
            break
        elif user_input.lower() == "n":
            print("Exiting...")
            exit()
        else:
            print("Invalid input. Please try again.")

    embeddings = OpenAIEmbeddings()
    if index_name not in pinecone.list_indexes():
        raise ValueError(f"Index {index_name} does not exist.")
    try:
        print(f"Indexing documents... ", end="")
        vector_store = Pinecone.from_documents(
            chunks, embeddings, index_name=index_name
        )
        print(f"Done")
    except pinecone.exceptions.PineconeException as e:
        raise pinecone.exceptions.PineconeException(
            f"Error indexing documents: {str(e)}"
        )
    return vector_store


def fetch_vector_store(index_name: str) -> Pinecone:
    """
    Retrieves embeddings from an already existing Pinecone index.

    Args:
        index_name (str): The name of the index to fetch embeddings from.

    Returns:
        Pinecone: A Pinecone vector store object containing the fetched embeddings.

    Raises:
        ValueError: If the index name does not exist.
        pinecone.exceptions.PineconeException: If there is an error fetching the embeddings.
    """
    import pinecone
    from langchain.vectorstores import Pinecone
    from langchain.embeddings.openai import OpenAIEmbeddings

    embeddings = OpenAIEmbeddings()
    if index_name not in pinecone.list_indexes():
        raise ValueError(f"Index {index_name} does not exist.")
    try:
        print(f"Fetching {index_name}... ", end="")
        vector_store = Pinecone.from_existing_index(
            index_name=index_name, embedding=embeddings
        )
        print(f"Done")
        stats = vector_store.get_pinecone_index(
            index_name=index_name
        ).describe_index_stats()
        print(
            f"Index contains {stats['total_vector_count']} vectors for as many chunks of data"
        )
    except pinecone.exceptions.PineconeException as e:
        raise pinecone.exceptions.PineconeException(
            f"Error fetching embeddings: {str(e)}"
        )
    return vector_store


def delete_pinecone_index(index_name: str) -> None:
    """
    Deletes a Pinecone index with the given name or all indices if index_name is "all".

    Args:
        index_name (str): Name of the index to delete. If "all", all indices will be deleted.

    Returns:
        None
    """
    import pinecone

    if index_name == "all":
        indices = pinecone.list_indexes()
        print(f"Deleting all {len(indices)} indices... ", end="")
        for index in indices:
            pinecone.delete_index(index)
        print("Done")
    else:
        print(f"Deleting index {index_name}... ", end="")
        pinecone.delete_index(index_name)
        print("Done")

pinecone_utils/pinecone_utils.py
bshastry/docubot
#!/usr/bin/env python3

import unittest


def run_tests():
    """
    Discovers and runs all tests in the 'tests' directory.

    This function uses the unittest module to discover and run all tests in the 'tests' directory.
    It creates a test loader to load the tests, a test suite to hold the discovered tests, and a test runner to execute the tests.
    The 'tests' directory should contain all the unit test files.

    Returns:
    None
    """

    test_loader = unittest.TestLoader()
    test_suite = test_loader.discover("tests")
    test_runner = unittest.TextTestRunner()
    test_runner.run(test_suite)


if __name__ == "__main__":
    run_tests()

test_runner.py
bshastry/docubot
import unittest
from docubot import build_kb
from text_utils.text_utils import tiktoken_len
from langchain.docstore.document import Document


class TestDocuBot(unittest.TestCase):
    """
    A class for testing the functionality of the DocuBot application.
    """

    def test_build_kb(self):
        chunks = build_kb("test_files")
        # Test if the function returns a list
        self.assertIsInstance(chunks, list)

        # Test if a chunk is a document
        self.assertIsInstance(chunks[0], Document)

        # Test if the function returns a list of chunks with at most 512 tokens per chunk
        for c in chunks:
            self.assertLessEqual(tiktoken_len(c.page_content), 512)


if __name__ == "__main__":
    unittest.main()

tests/test_docubot.py
bshastry/docubot
#!/usr/bin/env python3

"""
This module contains unit tests for the document loaders module.

The document loaders module provides functions for loading various types of documents, such as PDF, DOCX, Markdown, and plain text files. It also includes a function for loading content from Wikipedia.

The unit tests in this module ensure that the document loading functions return the expected results. They also test the chunk_data function, which is used to split document content into smaller chunks.

The module includes the following test cases:

- test_load_pdf_document: Tests the load_pdf_document function by loading a PDF file and checking that the returned content is a list with a length greater than 0.

- test_load_docx_document: Tests the load_docx_document function by loading a DOCX file and checking that the returned content is a list with a length greater than 0.

- test_load_markdown_document: Tests the load_markdown_document function by loading a Markdown file and checking that the returned content is a list with a length greater than 0.

- test_load_txt_document: Tests the load_txt_document function by loading a plain text file and checking that the returned content is a list with a length greater than 0.

- test_load_from_wikipedia: Tests the load_from_wikipedia function by loading content from Wikipedia for a given query and checking that the returned content is a list with a length greater than 0.

- test_load_document: Tests the load_document function by loading various types of documents (PDF, DOCX, Markdown, and plain text files) and checking that the returned content is a list with a length greater than 0.

- test_chunk_data: Tests the chunk_data function by loading a plain text file, splitting its content into smaller chunks, and checking that each chunk has a length less than or equal to the specified chunk size.

To run the unit tests, execute this module as a script.
"""

import unittest
from document_loaders.document_loaders import (
    load_pdf_document,
    load_docx_document,
    load_markdown_document,
    load_txt_document,
    load_from_wikipedia,
    load_document,
    chunk_data,
    merge_document,
)
from text_utils.text_utils import tiktoken_len
from typing import List, TypeVar

T = TypeVar("T")


class TestDocumentLoaders(unittest.TestCase):
    def test_load_pdf_document(self):
        pdf_file = "test_files/test.pdf"
        pdf_contents = load_pdf_document(pdf_file)
        self.assertIsInstance(pdf_contents, list)
        self.assertGreater(len(pdf_contents), 0)

    def test_load_docx_document(self):
        docx_file = "test_files/test.docx"
        docx_contents = load_docx_document(docx_file)
        self.assertIsInstance(docx_contents, list)
        self.assertGreater(len(docx_contents), 0)

    def test_load_markdown_document(self):
        md_file = "test_files/test.md"
        md_contents = load_markdown_document(md_file)
        self.assertIsInstance(md_contents, list)
        self.assertGreater(len(md_contents), 0)

    def test_load_txt_document(self):
        txt_file = "test_files/test.txt"
        txt_contents = load_txt_document(txt_file)
        self.assertIsInstance(txt_contents, list)
        self.assertGreater(len(txt_contents), 0)

    def test_load_from_wikipedia(self):
        query = "Python programming language"
        wikipedia_contents = load_from_wikipedia(query)
        self.assertIsInstance(wikipedia_contents, list)
        self.assertGreater(len(wikipedia_contents), 0)

    def test_load_document(self):
        pdf_file = "test_files/test.pdf"
        docx_file = "test_files/test.docx"
        md_file = "test_files/test.md"
        txt_file = "test_files/test.txt"
        url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf"

        pdf_contents = load_document(pdf_file)
        self.assertIsInstance(pdf_contents, list)
        self.assertGreater(len(pdf_contents), 0)

        docx_contents = load_document(docx_file)
        self.assertIsInstance(docx_contents, list)
        self.assertGreater(len(docx_contents), 0)

        md_contents = load_document(md_file)
        self.assertIsInstance(md_contents, list)
        self.assertGreater(len(md_contents), 0)

        txt_contents = load_document(txt_file)
        self.assertIsInstance(txt_contents, list)
        self.assertGreater(len(txt_contents), 0)

        url_contents = load_document(url)
        self.assertIsInstance(url_contents, list)

    def chunk_txt_file(self, chunk_size: int) -> List[T]:
        txt_file = "test_files/test.txt"
        txt_contents = load_document(txt_file)
        # Number of tokens to overlap between chunks
        chunk_overlap = 5
        chunks = chunk_data(
            txt_contents, chunk_size=chunk_size, chunk_overlap=chunk_overlap
        )
        return chunks

    def test_chunk_data(self):
        # Chunk size in tokens
        chunk_size = 10
        chunks = self.chunk_txt_file(chunk_size=chunk_size)
        self.assertIsInstance(chunks, list)
        self.assertGreater(len(chunks), 1)
        for chunk in chunks:
            self.assertLessEqual(tiktoken_len(chunk.page_content), chunk_size)

    def test_merge_single_document(self):
        txt_file = "test_files/test.txt"
        document = load_txt_document(txt_file)
        expected_output = "This is a text file that has more than ten characters.\n"
        self.assertEqual(merge_document(document), expected_output)

    def test_merge_multiple_documents(self):
        # Chunk size in tokens
        chunk_size = 10
        chunks = self.chunk_txt_file(chunk_size=chunk_size)
        expected_output = "This is a text file that has more than ten\n\nthat has more than ten characters."
        self.assertEqual(merge_document(chunks), expected_output)

    def test_merge_empty_document(self):
        document = []
        expected_output = ""
        self.assertEqual(merge_document(document), expected_output)


if __name__ == "__main__":
    unittest.main()

tests/test_document_loaders.py
bshastry/docubot
from text_utils.text_utils import (
    tiktoken_len,
    num_tokens_and_cost,
    return_url_extension,
)
import unittest

# As of 2021-10-20, the cost of embedding a single token using OpenAI is $0.0000001
EMBEDDING_COST_PER_TOKEN = 0.0000001


class TestTextUtils(unittest.TestCase):
    def test_tiktoken_len(self):
        self.assertEqual(tiktoken_len("Hello, world!"), 4)
        self.assertEqual(tiktoken_len("This is a sentence."), 5)
        self.assertEqual(tiktoken_len("This is a longer sentence with more words."), 9)

    def test_embedding_cost(self):
        class Page:
            def __init__(self, page_content):
                self.page_content = page_content

        document = [
            Page("This is the first page."),
            Page("This is the second page."),
            Page("This is the third page."),
        ]
        num_tokens, cost = num_tokens_and_cost(document)
        self.assertEqual(num_tokens, 18)
        self.assertAlmostEquals(cost, num_tokens * EMBEDDING_COST_PER_TOKEN)

        document = [
            Page("This is a short page."),
            Page("This is a longer page with more words."),
            Page("This is the longest page of them all, with many many words."),
        ]
        num_tokens, cost = num_tokens_and_cost(document)
        self.assertEqual(num_tokens, 29)
        self.assertAlmostEqual(cost, (num_tokens * EMBEDDING_COST_PER_TOKEN))

    def test_return_url_extension(self):
        self.assertEqual(
            return_url_extension("https://www.example.com/index.html"), ".html"
        )
        self.assertEqual(
            return_url_extension("https://www.example.com/path/to/file.txt"), ".txt"
        )
        self.assertEqual(
            return_url_extension("https://www.example.com/path/to/image.jpg"), ".jpg"
        )
        self.assertEqual(
            return_url_extension("https://www.example.com/path/to/doc.pdf"), ".pdf"
        )
        self.assertEqual(
            return_url_extension("https://www.example.com/path/to/README.md"), ".md"
        )
        self.assertEqual(
            return_url_extension("https://www.example.com/path/to/ms.docx"), ".docx"
        )


if __name__ == "__main__":
    unittest.main()

tests/test_text_utils.py
bshastry/docubot
#!/usr/bin/env python3

"""
This module provides utility functions for working with text data.

It includes functions for calculating the length of tokenized text using the TikToken library,
calculating the embedding cost for a list of texts, and extracting file extensions from URLs.

Note: This module requires the 'tiktoken' library to be installed.
"""
from typing import List, TypeVar

T = TypeVar("T")


# create the length function
def tiktoken_len(text: str) -> int:
    """
    Returns the length of the tokenized version of the input text using the TikToken library.

    Args:
        text (str): The input text to tokenize.

    Returns:
        int: The length of the tokenized version of the input text.
    """
    import tiktoken

    tokenizer = tiktoken.encoding_for_model("text-embedding-ada-002")
    tokens = tokenizer.encode(text, disallowed_special=())
    return len(tokens)


def num_tokens_and_cost(document: List[T]) -> (int, float):
    """
    Calculates the total number of tokens and cost for a list of texts.

    Args:
    - document (List[T]): A list of texts to calculate the number of tokens and cost for.

    Returns:
    - int: The total number of tokens for the given list of texts.
    - float: The cost for the given list of texts.
    """

    total_tokens = sum([tiktoken_len(page.page_content) for page in document])
    return total_tokens, (total_tokens / 1000) * 0.0001


def return_url_extension(url: str) -> str:
    """
    Return the file extension from the given URL.

    Args:
    url (str): The URL to extract the file extension from.

    Returns:
    str: The file extension from the given URL.
    """
    from urllib.parse import urlparse
    import os

    parsed = urlparse(url)
    _, ext = os.path.splitext(parsed.path)
    return ext

text_utils/text_utils.py
bshastry/docubot