This repository has been archived by the owner on Apr 13, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
/
vectordb.py
246 lines (202 loc) · 9.56 KB
/
vectordb.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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
from dotenv import load_dotenv
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.text_splitter import Language, RecursiveCharacterTextSplitter
from langchain.docstore.document import Document
import chromadb
import characters
import settings
import openai
import os
import json
import helper_module
from pathlib import Path
def load_document_single(filepath: str) -> Document:
# Loads a single document from a file path
file_extension = os.path.splitext(filepath)[1]
loader_class = settings.DOCUMENTS_MAP.get(file_extension)
if loader_class:
loader = loader_class(filepath)
else:
raise ValueError(f"Unknown document type: {filepath}")
return loader.load()[0]
def load_documents_batch(filepaths):
helper_module.log(f"Loading documents in batch: {filepaths}", 'info')
# create a thread pool
with ThreadPoolExecutor(len(filepaths)) as exe:
# load files
futures = [exe.submit(load_document_single, name) for name in filepaths]
# collect data
data_list = [future.result() for future in futures]
# return data and file paths
return data_list, filepaths
def split_documents(documents):
# Splits documents for correct Text Splitter
text_docs, python_docs = [], []
for doc in documents:
file_extension = os.path.splitext(doc.metadata["source"])[1]
if file_extension == ".py":
python_docs.append(doc)
else:
text_docs.append(doc)
return text_docs, python_docs
def load_documents(source_folder: str):
# Loads all documents from the source documents directory, including nested folders
paths = []
for root, _, files in os.walk(source_folder):
for file_name in files:
file_extension = os.path.splitext(file_name)[1]
source_file_path = os.path.join(root, file_name)
if file_extension in settings.DOCUMENTS_MAP.keys():
paths.append(source_file_path)
# Have at least one worker and at most INGEST_THREADS workers
n_workers = min(settings.INGEST_THREADS, max(len(paths), 1))
chunk_size = round(len(paths) / n_workers)
docs = []
with ProcessPoolExecutor(n_workers) as executor:
futures = []
# split the load operations into chunks
for i in range(0, len(paths), chunk_size):
# select a chunk of filenames
filepaths = paths[i : (i + chunk_size)]
# submit the task
future = executor.submit(load_documents_batch, filepaths)
futures.append(future)
# process all results
for future in as_completed(futures):
# open the file and load the data
contents, _ = future.result()
docs.extend(contents)
return docs
def create_vectordb(source_folder):
if not os.path.exists(source_folder):
os.makedirs(source_folder)
# Load documents and split in chunks
helper_module.log(f"Loading documents from {source_folder}", 'info')
documents = load_documents(source_folder)
text_documents, python_documents = split_documents(documents)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=settings.DOCUMENT_SPLITTER_CHUNK_SIZE,
chunk_overlap=settings.DOCUMENT_SPLITTER_CHUNK_OVERLAP)
python_splitter = RecursiveCharacterTextSplitter.from_language(
language=Language.PYTHON,
chunk_size=settings.PYTHON_SPLITTER_CHUNK_SIZE,
chunk_overlap=settings.PYTHON_SPLITTER_CHUNK_OVERLAP
)
texts = text_splitter.split_documents(text_documents)
texts.extend(python_splitter.split_documents(python_documents))
helper_module.log(f"Loaded {len(documents)} documents from {source_folder}", 'info')
helper_module.log(f"{len(texts)} chunks of text split", 'info')
embeddings = OpenAIEmbeddings()
my_vectordb = Chroma.from_documents(
collection_name=settings.VECTORDB_COLLECTION,
documents=texts,
embedding=embeddings,
persist_directory=settings.CHROMA_DB_FOLDER,
)
return my_vectordb
def get_vectordb(collection_name=settings.VECTORDB_COLLECTION):
embeddings = OpenAIEmbeddings()
my_vectordb = Chroma(
collection_name=collection_name,
persist_directory=settings.CHROMA_DB_FOLDER,
embedding_function=embeddings
)
return my_vectordb
def delete_vectordb():
my_vectordb = get_vectordb(settings.VECTORDB_COLLECTION)
my_vectordb.delete_collection()
my_vectordb.persist()
helper_module.log(f"Vector DB collection deleted: {settings.VECTORDB_COLLECTION}", 'info')
def retrieval_qa_run(system_message, human_input, context_memory, callbacks=None):
my_vectordb = get_vectordb()
retriever = my_vectordb.as_retriever(search_kwargs={"k": settings.NUM_SOURCES_TO_RETURN})
template = system_message + settings.RETRIEVER_TEMPLATE
qa_prompt = PromptTemplate(input_variables=["history", "context", "question"],
template=template)
qa_chain = RetrievalQA.from_chain_type(
llm=ChatOpenAI(
temperature=settings.DEFAULT_GPT_QA_HELPER_MODEL_TEMPERATURE,
model_name=settings.DEFAULT_GPT_QA_HELPER_MODEL,
streaming=True,
callbacks=callbacks,
),
chain_type="stuff",
retriever=retriever,
return_source_documents=True,
chain_type_kwargs={"prompt": qa_prompt,
"memory": context_memory},
)
helper_module.log("Running QA chain...", 'info')
response = qa_chain(human_input)
my_answer, my_docs = response["result"], response["source_documents"]
helper_module.log(f"Answer: {my_answer}", 'info')
return my_answer, my_docs
def embed_conversations():
""" Ingest past conversations as long-term memory into the vector DB."""
helper_module.log(f"Loading conversations in batch: {settings.CONVERSATION_SAVE_FOLDER}", 'info')
conversations = []
for json_file in settings.CONVERSATION_SAVE_FOLDER.glob('*.json'):
if not str(json_file).endswith(settings.SNAPSHOT_FILENAME):
with open(json_file, 'r') as f:
json_data = json.load(f)
result_str = ""
for entry in json_data:
result_str += f"{entry['role']}: {entry['content']}\n"
conversations.append(Document(page_content=result_str, metadata = {"source": str(json_file)}))
text_splitter = RecursiveCharacterTextSplitter(chunk_size=settings.DOCUMENT_SPLITTER_CHUNK_SIZE,
chunk_overlap=settings.DOCUMENT_SPLITTER_CHUNK_OVERLAP)
texts = text_splitter.split_documents(conversations)
my_vectordb = get_vectordb()
my_vectordb.add_documents(documents=texts, embeddings=OpenAIEmbeddings())
helper_module.log(f"{len(conversations)} of conversations found", 'info')
helper_module.log(f"{len(texts)} chunks of text embedded", 'info')
def display_vectordb_info():
persistent_client = chromadb.PersistentClient(path=settings.CHROMA_DB_FOLDER)
collection = persistent_client.get_or_create_collection(settings.VECTORDB_COLLECTION)
helper_module.log(f"VectorDB Folder: {settings.CONVERSATION_SAVE_FOLDER}", 'info')
helper_module.log(f"Collection: {settings.VECTORDB_COLLECTION}", 'info')
helper_module.log(f"Number of items in collection: {collection.count()}", 'info')
if __name__ == "__main__":
load_dotenv()
openai.api_key = os.getenv('OPENAI_API_KEY')
while True:
display_vectordb_info()
user_input = input("\n(C)create DB, (E)mbed conversations, (D)elete collection, (Q)uery - type 'quit' or 'exit' to quit: ")
if 'c' == user_input.lower().strip():
create_vectordb(settings.DOCUMENT_FOLDER)
elif 'e' == user_input.lower().strip():
embed_conversations()
elif 'd' == user_input.lower().strip():
user_input = input("\nAre you sure? Type 'yes' if you are: ")
if 'yes' == user_input.lower().strip():
delete_vectordb()
elif 'q' == user_input.lower().strip():
while True:
memory = ConversationBufferMemory(input_key="question",
memory_key="history")
query = input("\nQuery: ")
if 'quit' in query or 'exit' in query:
break
helper_module.log(f"Querying model: {settings.DEFAULT_GPT_QA_HELPER_MODEL}", 'info')
system_input = characters.CUSTOM_INSTRUCTIONS
answer, docs = retrieval_qa_run(system_input, query, memory)
# Print the result
print("\n\n> Question:")
print(query)
print("\n> Answer:")
print(answer)
if settings.SHOW_SOURCES:
print("----------------------------------SOURCE DOCUMENTS---------------------------")
for document in docs:
print("\n> " + document.metadata["source"] + ":")
print(document.page_content)
print("----------------------------------SOURCE DOCUMENTS---------------------------")
elif 'quit' in user_input or 'exit' in user_input:
break
else:
print("Unknown choice.\n")