solicit.py

""" Solicit John's Brain with a Question """

import os
from llama_index import GPTSimpleVectorIndex

# noinspection SpellCheckingInspection
os.environ["OPENAI_API_KEY"] = "YOUR KEY"


def load_index(index_location: str):
    """
    Load an GPTSimpleVectorIndex from a filename

    :param index_location: The name of the file you want to load
    :return:
    """
    return GPTSimpleVectorIndex.load_from_disk(index_location)


def query_index_file(index_file: str, query_string: str, context: int = 20):
    """
    Send a ChatGPT query using the context generated by index_file.

    :param index_file: The filename of a previously saved GPTSimpleVectorIndex
    :param query_string: The string of a question you want to ask
    :param context: How much "context" should be provided to ChatGPT? More gives better results but is slower and
        more costly
    :return: A response object with the answer
    """
    index = load_index(index_file)
    return index.query(query_string, similarity_top_k=context)


if __name__ == "__main__":
    import sys

    if not sys.argv[1]:
        print('Usage: python solicit.py "My question"')

    print(query_index_file("obsidian_index.json", sys.argv[1]))