/
llama_index_buddy.py
33 lines (28 loc) · 968 Bytes
/
llama_index_buddy.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
### inspired by the hello world example
import sys
import os.path
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
directory = sys.argv[1]
query = sys.argv[2]
from llama_index.core import (
VectorStoreIndex,
SimpleDirectoryReader,
StorageContext,
load_index_from_storage,
)
# check if index already exists
PERSIST_DIR = "./" + directory + "_index"
if not os.path.exists(PERSIST_DIR):
# load the documents and create the index
documents = SimpleDirectoryReader(directory).load_data()
index = VectorStoreIndex.from_documents(documents)
# store it for later
index.storage_context.persist(persist_dir=PERSIST_DIR)
else:
# load the existing index
storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
index = load_index_from_storage(storage_context)
# Either way we can now query the index
query_engine = index.as_query_engine()
response = query_engine.query(query)
print(response)