Update rag.py

app/rag.py

@@ -24,6 +24,9 @@ logger = logging.getLogger(__name__)
 
 class ChatPDF:
     pdf_count = 0
+    text_chunks = []
+    doc_ids = []
+    nodes = []
 
     def __init__(self):
         self.text_parser = SentenceSplitter(chunk_size=1024, chunk_overlap=24)
@@ -63,28 +66,24 @@ class ChatPDF:
         Settings.transformations = [self.text_parser]
 
     def ingest(self, files_dir: str):
-        text_chunks = []
-        doc_ids = []
-        nodes = []
-
         docs = SimpleDirectoryReader(input_dir=files_dir).load_data()
 
         logger.info("enumerating docs")
         for doc_idx, doc in enumerate(docs):
             self.pdf_count = self.pdf_count + 1
             curr_text_chunks = self.text_parser.split_text(doc.text)
-            text_chunks.extend(curr_text_chunks)
-            doc_ids.extend([doc_idx] * len(curr_text_chunks))
+            self.text_chunks.extend(curr_text_chunks)
+            self.doc_ids.extend([doc_idx] * len(curr_text_chunks))
 
         logger.info("enumerating text_chunks")
-        for idx, text_chunk in enumerate(text_chunks):
+        for idx, text_chunk in enumerate(self.text_chunks):
             node = TextNode(text=text_chunk)
-            src_doc = docs[doc_ids[idx]]
+            self.src_doc = docs[doc_ids[idx]]
             node.metadata = src_doc.metadata
-            nodes.append(node)
+            self.nodes.append(node)
 
         logger.info("enumerating nodes")
-        for node in nodes:
+        for node in self.nodes:
             node_embedding = self.embed_model.get_text_embedding(
                 node.get_content(metadata_mode=MetadataMode.ALL)
             )
@@ -94,7 +93,7 @@ class ChatPDF:
         storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
         logger.info("indexing the nodes in VectorStoreIndex")
         index = VectorStoreIndex(
-            nodes=nodes,
+            nodes=self.nodes,
             storage_context=storage_context,
             transformations=Settings.transformations,
         )
@@ -110,5 +109,8 @@ class ChatPDF:
         return streaming_response
 
     def clear(self):
+        self.self.client.delete(collection_name="rag_documents", ids=doc_ids)
         self.pdf_count = 0
-        self.client.delete(collection_name="rag_documents")
+        self.text_chunks = []
+        self.doc_ids = []
+        self.nodes = []