fix persistence

hello_wordsmith/wordsmith.py

@@ -2,43 +2,38 @@ import os
 import sys
 
 import chromadb
-from llama_index.cli.rag import RagCLI
-from llama_index.core import (
-    ChatPromptTemplate,
-    SimpleDirectoryReader,
-    VectorStoreIndex
-)
+from llama_index.cli.rag import RagCLI, default_ragcli_persist_dir
+from llama_index.core import (ChatPromptTemplate, Settings,
+                              SimpleDirectoryReader, StorageContext,
+                              VectorStoreIndex)
 from llama_index.core.base.llms.types import ChatMessage, MessageRole
-from llama_index.core.ingestion import IngestionPipeline
+from llama_index.core.ingestion import IngestionCache, IngestionPipeline
 from llama_index.core.query_pipeline import InputComponent, QueryPipeline
 from llama_index.core.response_synthesizers import TreeSummarize
-from llama_index.core import Settings
-from llama_index.embeddings.openai import (
-    OpenAIEmbedding,
-    OpenAIEmbeddingModelType
-)
+from llama_index.core.storage.docstore import SimpleDocumentStore
+from llama_index.embeddings.openai import (OpenAIEmbedding,
+                                           OpenAIEmbeddingModelType)
 from llama_index.llms.openai import OpenAI
 from llama_index.vector_stores.chroma import ChromaVectorStore
 
-
-Settings.embed_model = OpenAIEmbedding(
-    model="text-embedding-3-small"
-)
+Settings.embed_model = OpenAIEmbedding(model=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL)
 
 
 def initialize_chroma_db():
-    chroma_client = chromadb.EphemeralClient()
-    chroma_collection = chroma_client.create_collection("wordsmith")
+    db = chromadb.PersistentClient(
+        path=os.path.join(default_ragcli_persist_dir(), "chroma")
+    )
+    chroma_collection = db.get_or_create_collection("wordsmith_rag_demo_index")
     vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
     return vector_store
 
 
-def setup_document_storage(vector_store):
+def setup_document_storage(*, vector_store, storage_context):
     package_directory = os.path.dirname(os.path.abspath(__file__))
     dataset_path = os.path.join(package_directory, "public_wordsmith_dataset")
     reader = SimpleDirectoryReader(input_dir=dataset_path)
     docs = reader.load_data()
-    index = VectorStoreIndex.from_documents(docs)
+    index = VectorStoreIndex.from_documents(docs, storage_context=storage_context)
     return index
 
 
@@ -84,9 +79,7 @@ _chat_template_messages = [
 
 def configure_query_pipeline(index, llm):
     """Configure and set up the query pipeline"""
-    text_qa_chat_template = ChatPromptTemplate.from_messages(
-        _chat_template_messages
-    )
+    text_qa_chat_template = ChatPromptTemplate.from_messages(_chat_template_messages)
     query_pipeline = QueryPipeline()
 
     retriever = index.as_retriever(similarity_top_k=20)
@@ -122,14 +115,19 @@ class WordsmithRAGCLI(RagCLI):
 
 def main():
     vector_store = initialize_chroma_db()
-    index = setup_document_storage(vector_store)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    index = setup_document_storage(
+        vector_store=vector_store, storage_context=storage_context
+    )
     llm = initialize_llm()
     query_pipeline = configure_query_pipeline(index, llm)
-    ingestion_pipeline = IngestionPipeline(vector_store=vector_store)
+    ingestion_pipeline = IngestionPipeline(
+        vector_store=vector_store,
+        cache=IngestionCache(),
+        docstore=SimpleDocumentStore(),
+    )
     rag_cli_instance = WordsmithRAGCLI(
-        ingestion_pipeline=ingestion_pipeline,
-        llm=llm,
-        query_pipeline=query_pipeline
+        ingestion_pipeline=ingestion_pipeline, llm=llm, query_pipeline=query_pipeline
     )
     rag_cli_instance.cli()