-
Notifications
You must be signed in to change notification settings - Fork 13
/
index_manager.py
90 lines (61 loc) · 1.95 KB
/
index_manager.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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
from llama_index.core import (
VectorStoreIndex,
StorageContext,
ServiceContext,
SimpleDirectoryReader
)
from llama_index.core.indices.loading import load_index_from_storage
from utils.vector_database import build_pinecone_vector_store, build_mongo_index
from mongodb.index import getExistingLlamaIndexes
from llama_index.llms.perplexity import Perplexity
import os
llm = Perplexity(
api_key=os.getenv("PERPLEXITY_API_KEY"), model="mixtral-8x7b-instruct", temperature=0.2
)
index_store = build_mongo_index()
vector_store = build_pinecone_vector_store()
service_context = ServiceContext.from_defaults(
llm=llm,
embed_model="local:BAAI/bge-small-en-v1.5",
)
storage_context = StorageContext.from_defaults(
index_store=index_store,
vector_store=vector_store
)
mongoIndex = None
def initialize_index():
existing_indexes = getExistingLlamaIndexes()
global mongoIndex
if len(existing_indexes) > 0:
print("Loading existing index...")
mongoIndex = load_index_from_storage(
service_context=service_context,
storage_context=storage_context,
llm=llm,
index_id='mongo-index',
)
return createQueryEngine(mongoIndex)
else:
print("Building index...")
mongoIndex = buildVectorIndex()
return createQueryEngine(mongoIndex)
def createQueryEngine(index):
return index.as_retriever(top_k=3)
def get_service_context():
return service_context
def update_index(doc):
mongoIndex.insert(doc)
def delete_document_from_index(doc_id):
mongoIndex.delete_ref_doc(doc_id)
def buildVectorIndex():
reader = SimpleDirectoryReader(
input_files=["./data/rules.pdf"]
)
documents = reader.load_data()
index = VectorStoreIndex.from_documents(
documents,
service_context=service_context,
storage_context=storage_context
)
index.set_index_id("mongo-index")
return index