break apart into modules

hello_wordsmith/datastores.py

@@ -0,0 +1,65 @@
+import os
+
+import chromadb
+from llama_index.cli.rag import default_ragcli_persist_dir
+from llama_index.core import SimpleDirectoryReader, StorageContext, VectorStoreIndex
+from llama_index.core.storage.docstore import SimpleDocumentStore
+from llama_index.vector_stores.chroma import ChromaVectorStore
+from pydantic.v1 import BaseModel
+
+
+class InitialisedDataContainer(BaseModel):
+    class Config:
+        arbitrary_types_allowed = True
+
+    db: chromadb.Collection
+    doc_store: SimpleDocumentStore
+    vector_store: ChromaVectorStore
+    index: VectorStoreIndex
+    storage_context: StorageContext
+
+
+def _get_chroma_db() -> chromadb.Collection:
+    db = chromadb.PersistentClient(
+        path=os.path.join(default_ragcli_persist_dir(), "chroma")
+    )
+    chroma_collection = db.get_or_create_collection("wordsmith_rag_demo_index")
+    return chroma_collection
+
+
+def fetch_or_initialise_datastores() -> InitialisedDataContainer:
+    db = _get_chroma_db()
+    vector_store = ChromaVectorStore(chroma_collection=db)
+    try:
+        docstore = SimpleDocumentStore.from_persist_dir(
+            persist_dir=os.path.join(default_ragcli_persist_dir(), "storage")
+        )
+    except FileNotFoundError:
+        docstore = SimpleDocumentStore()
+    storage_context = StorageContext.from_defaults(
+        vector_store=vector_store, docstore=docstore
+    )
+    if not docstore.docs or not db.count():
+        package_directory = os.path.dirname(os.path.abspath(__file__))
+        dataset_path = os.path.join(package_directory, "public_wordsmith_dataset")
+        docs = SimpleDirectoryReader(
+            input_dir=dataset_path, filename_as_id=True
+        ).load_data()
+        docstore.add_documents(docs)
+        docstore.persist(
+            persist_path=os.path.join(
+                default_ragcli_persist_dir(), "./storage/docstore.json"
+            )
+        )
+        index = VectorStoreIndex.from_documents(
+            documents=docs, storage_context=storage_context
+        )
+    else:
+        index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+    return InitialisedDataContainer(
+        db=db,
+        doc_store=docstore,
+        vector_store=vector_store,
+        index=index,
+        storage_context=storage_context,
+    )

hello_wordsmith/query_pipeline.py

@@ -0,0 +1,65 @@
+from llama_index.core import ChatPromptTemplate, VectorStoreIndex
+from llama_index.core.base.llms.types import ChatMessage, MessageRole
+from llama_index.core.query_pipeline import InputComponent, QueryPipeline
+from llama_index.core.response_synthesizers import TreeSummarize
+from llama_index.llms.openai import OpenAI
+
+_system_prompt = ChatMessage(
+    content=(
+        "You are an expert Q&A analyst representing Wordsmith in front of "
+        "potentially interested users.\n"
+        "If the question is related to Wordsmith in any way, "
+        "answer the query using the provided context information.\n"
+        "If you can't find the answer in the provided context information, "
+        "simply say you don't have enough information to answer the query.\n"
+        "Always be polite and professional.\n"
+        "Some rules to follow:\n"
+        "1. Never directly reference the given context in your answer.\n"
+        "2. Avoid statements like 'Based on the context, ...' or "
+        "'The context information ...', etc."
+    ),
+    role=MessageRole.SYSTEM,
+)
+
+_chat_template_messages = [
+    _system_prompt,
+    ChatMessage(
+        content=(
+            "Context information from multiple sources is below.\n"
+            "---------------------\n"
+            "{context_str}\n"
+            "---------------------\n"
+            "Given the information from multiple sources and not prior knowledge, "
+            "answer the query.\n"
+            "Query: {query_str}\n"
+            "Answer: "
+        ),
+        role=MessageRole.USER,
+    ),
+]
+
+_TOP_K_RETRIEVAL = 20
+
+
+def configure_query_pipeline(*, index: VectorStoreIndex, llm: OpenAI) -> QueryPipeline:
+    """Configure and set up the query pipeline"""
+    text_qa_chat_template = ChatPromptTemplate.from_messages(_chat_template_messages)
+    query_pipeline = QueryPipeline()
+
+    retriever = index.as_retriever(similarity_top_k=_TOP_K_RETRIEVAL)
+    summarizer = TreeSummarize(
+        llm=llm, streaming=True, summary_template=text_qa_chat_template
+    )
+
+    query_pipeline.add_modules(
+        {
+            "input": InputComponent(),
+            "retriever": retriever,
+            "summarizer": summarizer,
+        }
+    )
+    query_pipeline.add_link("input", "retriever")
+    query_pipeline.add_link("input", "summarizer", dest_key="query_str")
+    query_pipeline.add_link("retriever", "summarizer", dest_key="nodes")
+
+    return query_pipeline

hello_wordsmith/wordsmith.py

@@ -1,104 +1,15 @@
 import os
 import sys
+from typing import Callable
 
-import chromadb
-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.cli.rag import RagCLI
+from llama_index.core import Settings
 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.storage.docstore import SimpleDocumentStore
-from llama_index.embeddings.openai import (OpenAIEmbedding,
-                                           OpenAIEmbeddingModelType)
+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=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL)
-
-
-def initialize_chroma_db():
-    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, 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, storage_context=storage_context)
-    return index
-
-
-def initialize_llm():
-    llm = OpenAI(api_key=os.environ["OPENAI_API_KEY"], model="gpt-4")
-    return llm
-
-
-_system_prompt = ChatMessage(
-    content=(
-        "You are an expert Q&A analyst representing Wordsmith in front of "
-        "potentially interested users.\n"
-        "If the question is related to Wordsmith in any way, "
-        "answer the query using the provided context information.\n"
-        "If you can't find the answer in the provided context information, "
-        "simply say you don't have enough information to answer the query.\n"
-        "Always be polite and professional.\n"
-        "Some rules to follow:\n"
-        "1. Never directly reference the given context in your answer.\n"
-        "2. Avoid statements like 'Based on the context, ...' or "
-        "'The context information ...', etc."
-    ),
-    role=MessageRole.SYSTEM,
-)
-
-_chat_template_messages = [
-    _system_prompt,
-    ChatMessage(
-        content=(
-            "Context information from multiple sources is below.\n"
-            "---------------------\n"
-            "{context_str}\n"
-            "---------------------\n"
-            "Given the information from multiple sources and not prior knowledge, "
-            "answer the query.\n"
-            "Query: {query_str}\n"
-            "Answer: "
-        ),
-        role=MessageRole.USER,
-    ),
-]
-
-
-def configure_query_pipeline(index, llm):
-    """Configure and set up the query pipeline"""
-    text_qa_chat_template = ChatPromptTemplate.from_messages(_chat_template_messages)
-    query_pipeline = QueryPipeline()
-
-    retriever = index.as_retriever(similarity_top_k=20)
-    summarizer = TreeSummarize(
-        llm=llm, streaming=True, summary_template=text_qa_chat_template
-    )
-
-    query_pipeline.add_modules(
-        {
-            "input": InputComponent(),
-            "retriever": retriever,
-            "summarizer": summarizer,
-        }
-    )
-    query_pipeline.add_link("input", "retriever")
-    query_pipeline.add_link("input", "summarizer", dest_key="query_str")
-    query_pipeline.add_link("retriever", "summarizer", dest_key="nodes")
-
-    return query_pipeline
+from .datastores import fetch_or_initialise_datastores
+from .query_pipeline import configure_query_pipeline
 
 
 class WordsmithRAGCLI(RagCLI):
@@ -113,22 +24,31 @@ class WordsmithRAGCLI(RagCLI):
         super().cli()
 
 
-def main():
-    api_key = os.getenv("OPENAI_API_KEY")
-    if not api_key:
-        print("Error: Environment variable 'OPENAI_API_KEY' is not set. Please set this before running.")
-        sys.exit(1)
-    vector_store = initialize_chroma_db()
-    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)
+def _init_env(func: Callable[[], None]) -> Callable[[], None]:
+    def wrapper() -> None:
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            print(
+                "Error: Environment variable 'OPENAI_API_KEY' is not set. Please set this before running."
+            )
+            sys.exit(1)
+        Settings.embed_model = OpenAIEmbedding(
+            model=OpenAIEmbeddingModelType.TEXT_EMBED_3_SMALL
+        )
+        return func()
+
+    return wrapper
+
+
+@_init_env
+def main() -> None:
+    datastore_container = fetch_or_initialise_datastores()
+    llm = OpenAI(api_key=os.environ["OPENAI_API_KEY"], model="gpt-4")
+    query_pipeline = configure_query_pipeline(index=datastore_container.index, llm=llm)
     ingestion_pipeline = IngestionPipeline(
-        vector_store=vector_store,
+        vector_store=datastore_container.vector_store,
         cache=IngestionCache(),
-        docstore=SimpleDocumentStore(),
+        docstore=datastore_container.doc_store,
     )
     rag_cli_instance = WordsmithRAGCLI(
         ingestion_pipeline=ingestion_pipeline, llm=llm, query_pipeline=query_pipeline

poetry.lock

@@ -265,6 +265,52 @@ charset-normalizer = ["charset-normalizer"]
 html5lib = ["html5lib"]
 lxml = ["lxml"]
 
+[[package]]
+name = "black"
+version = "24.4.2"
+description = "The uncompromising code formatter."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "black-24.4.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:dd1b5a14e417189db4c7b64a6540f31730713d173f0b63e55fabd52d61d8fdce"},
+    {file = "black-24.4.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e537d281831ad0e71007dcdcbe50a71470b978c453fa41ce77186bbe0ed6021"},
+    {file = "black-24.4.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eaea3008c281f1038edb473c1aa8ed8143a5535ff18f978a318f10302b254063"},
+    {file = "black-24.4.2-cp310-cp310-win_amd64.whl", hash = "sha256:7768a0dbf16a39aa5e9a3ded568bb545c8c2727396d063bbaf847df05b08cd96"},
+    {file = "black-24.4.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:257d724c2c9b1660f353b36c802ccece186a30accc7742c176d29c146df6e474"},
+    {file = "black-24.4.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bdde6f877a18f24844e381d45e9947a49e97933573ac9d4345399be37621e26c"},
+    {file = "black-24.4.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e151054aa00bad1f4e1f04919542885f89f5f7d086b8a59e5000e6c616896ffb"},
+    {file = "black-24.4.2-cp311-cp311-win_amd64.whl", hash = "sha256:7e122b1c4fb252fd85df3ca93578732b4749d9be076593076ef4d07a0233c3e1"},
+    {file = "black-24.4.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:accf49e151c8ed2c0cdc528691838afd217c50412534e876a19270fea1e28e2d"},
+    {file = "black-24.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:88c57dc656038f1ab9f92b3eb5335ee9b021412feaa46330d5eba4e51fe49b04"},
+    {file = "black-24.4.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:be8bef99eb46d5021bf053114442914baeb3649a89dc5f3a555c88737e5e98fc"},
+    {file = "black-24.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:415e686e87dbbe6f4cd5ef0fbf764af7b89f9057b97c908742b6008cc554b9c0"},
+    {file = "black-24.4.2-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:bf10f7310db693bb62692609b397e8d67257c55f949abde4c67f9cc574492cc7"},
+    {file = "black-24.4.2-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:98e123f1d5cfd42f886624d84464f7756f60ff6eab89ae845210631714f6db94"},
+    {file = "black-24.4.2-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:48a85f2cb5e6799a9ef05347b476cce6c182d6c71ee36925a6c194d074336ef8"},
+    {file = "black-24.4.2-cp38-cp38-win_amd64.whl", hash = "sha256:b1530ae42e9d6d5b670a34db49a94115a64596bc77710b1d05e9801e62ca0a7c"},
+    {file = "black-24.4.2-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:37aae07b029fa0174d39daf02748b379399b909652a806e5708199bd93899da1"},
+    {file = "black-24.4.2-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:da33a1a5e49c4122ccdfd56cd021ff1ebc4a1ec4e2d01594fef9b6f267a9e741"},
+    {file = "black-24.4.2-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ef703f83fc32e131e9bcc0a5094cfe85599e7109f896fe8bc96cc402f3eb4b6e"},
+    {file = "black-24.4.2-cp39-cp39-win_amd64.whl", hash = "sha256:b9176b9832e84308818a99a561e90aa479e73c523b3f77afd07913380ae2eab7"},
+    {file = "black-24.4.2-py3-none-any.whl", hash = "sha256:d36ed1124bb81b32f8614555b34cc4259c3fbc7eec17870e8ff8ded335b58d8c"},
+    {file = "black-24.4.2.tar.gz", hash = "sha256:c872b53057f000085da66a19c55d68f6f8ddcac2642392ad3a355878406fbd4d"},
+]
+
+[package.dependencies]
+click = ">=8.0.0"
+mypy-extensions = ">=0.4.3"
+packaging = ">=22.0"
+pathspec = ">=0.9.0"
+platformdirs = ">=2"
+tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
+typing-extensions = {version = ">=4.0.1", markers = "python_version < \"3.11\""}
+
+[package.extras]
+colorama = ["colorama (>=0.4.3)"]
+d = ["aiohttp (>=3.7.4)", "aiohttp (>=3.7.4,!=3.9.0)"]
+jupyter = ["ipython (>=7.8.0)", "tokenize-rt (>=3.2.0)"]
+uvloop = ["uvloop (>=0.15.2)"]
+
 [[package]]
 name = "build"
 version = "1.2.1"
@@ -1594,6 +1640,53 @@ files = [
     {file = "multidict-6.0.5.tar.gz", hash = "sha256:f7e301075edaf50500f0b341543c41194d8df3ae5caf4702f2095f3ca73dd8da"},
 ]
 
+[[package]]
+name = "mypy"
+version = "1.10.0"
+description = "Optional static typing for Python"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "mypy-1.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:da1cbf08fb3b851ab3b9523a884c232774008267b1f83371ace57f412fe308c2"},
+    {file = "mypy-1.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:12b6bfc1b1a66095ab413160a6e520e1dc076a28f3e22f7fb25ba3b000b4ef99"},
+    {file = "mypy-1.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e36fb078cce9904c7989b9693e41cb9711e0600139ce3970c6ef814b6ebc2b2"},
+    {file = "mypy-1.10.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:2b0695d605ddcd3eb2f736cd8b4e388288c21e7de85001e9f85df9187f2b50f9"},
+    {file = "mypy-1.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:cd777b780312ddb135bceb9bc8722a73ec95e042f911cc279e2ec3c667076051"},
+    {file = "mypy-1.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3be66771aa5c97602f382230165b856c231d1277c511c9a8dd058be4784472e1"},
+    {file = "mypy-1.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8b2cbaca148d0754a54d44121b5825ae71868c7592a53b7292eeb0f3fdae95ee"},
+    {file = "mypy-1.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1ec404a7cbe9fc0e92cb0e67f55ce0c025014e26d33e54d9e506a0f2d07fe5de"},
+    {file = "mypy-1.10.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e22e1527dc3d4aa94311d246b59e47f6455b8729f4968765ac1eacf9a4760bc7"},
+    {file = "mypy-1.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:a87dbfa85971e8d59c9cc1fcf534efe664d8949e4c0b6b44e8ca548e746a8d53"},
+    {file = "mypy-1.10.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:a781f6ad4bab20eef8b65174a57e5203f4be627b46291f4589879bf4e257b97b"},
+    {file = "mypy-1.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b808e12113505b97d9023b0b5e0c0705a90571c6feefc6f215c1df9381256e30"},
+    {file = "mypy-1.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f55583b12156c399dce2df7d16f8a5095291354f1e839c252ec6c0611e86e2e"},
+    {file = "mypy-1.10.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4cf18f9d0efa1b16478c4c129eabec36148032575391095f73cae2e722fcf9d5"},
+    {file = "mypy-1.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:bc6ac273b23c6b82da3bb25f4136c4fd42665f17f2cd850771cb600bdd2ebeda"},
+    {file = "mypy-1.10.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:9fd50226364cd2737351c79807775136b0abe084433b55b2e29181a4c3c878c0"},
+    {file = "mypy-1.10.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:f90cff89eea89273727d8783fef5d4a934be2fdca11b47def50cf5d311aff727"},
+    {file = "mypy-1.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fcfc70599efde5c67862a07a1aaf50e55bce629ace26bb19dc17cece5dd31ca4"},
+    {file = "mypy-1.10.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:075cbf81f3e134eadaf247de187bd604748171d6b79736fa9b6c9685b4083061"},
+    {file = "mypy-1.10.0-cp38-cp38-win_amd64.whl", hash = "sha256:3f298531bca95ff615b6e9f2fc0333aae27fa48052903a0ac90215021cdcfa4f"},
+    {file = "mypy-1.10.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:fa7ef5244615a2523b56c034becde4e9e3f9b034854c93639adb667ec9ec2976"},
+    {file = "mypy-1.10.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:3236a4c8f535a0631f85f5fcdffba71c7feeef76a6002fcba7c1a8e57c8be1ec"},
+    {file = "mypy-1.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4a2b5cdbb5dd35aa08ea9114436e0d79aceb2f38e32c21684dcf8e24e1e92821"},
+    {file = "mypy-1.10.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:92f93b21c0fe73dc00abf91022234c79d793318b8a96faac147cd579c1671746"},
+    {file = "mypy-1.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:28d0e038361b45f099cc086d9dd99c15ff14d0188f44ac883010e172ce86c38a"},
+    {file = "mypy-1.10.0-py3-none-any.whl", hash = "sha256:f8c083976eb530019175aabadb60921e73b4f45736760826aa1689dda8208aee"},
+    {file = "mypy-1.10.0.tar.gz", hash = "sha256:3d087fcbec056c4ee34974da493a826ce316947485cef3901f511848e687c131"},
+]
+
+[package.dependencies]
+mypy-extensions = ">=1.0.0"
+tomli = {version = ">=1.1.0", markers = "python_version < \"3.11\""}
+typing-extensions = ">=4.1.0"
+
+[package.extras]
+dmypy = ["psutil (>=4.0)"]
+install-types = ["pip"]
+mypyc = ["setuptools (>=50)"]
+reports = ["lxml"]
+
 [[package]]
 name = "mypy-extensions"
 version = "1.0.0"
@@ -2053,8 +2146,8 @@ files = [
 [package.dependencies]
 numpy = [
     {version = ">=1.20.3", markers = "python_version < \"3.10\""},
-    {version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
     {version = ">=1.23.2", markers = "python_version >= \"3.11\""},
+    {version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
 ]
 python-dateutil = ">=2.8.2"
 pytz = ">=2020.1"
@@ -2083,6 +2176,17 @@ sql-other = ["SQLAlchemy (>=1.4.16)"]
 test = ["hypothesis (>=6.34.2)", "pytest (>=7.3.2)", "pytest-asyncio (>=0.17.0)", "pytest-xdist (>=2.2.0)"]
 xml = ["lxml (>=4.6.3)"]
 
+[[package]]
+name = "pathspec"
+version = "0.12.1"
+description = "Utility library for gitignore style pattern matching of file paths."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "pathspec-0.12.1-py3-none-any.whl", hash = "sha256:a0d503e138a4c123b27490a4f7beda6a01c6f288df0e4a8b79c7eb0dc7b4cc08"},
+    {file = "pathspec-0.12.1.tar.gz", hash = "sha256:a482d51503a1ab33b1c67a6c3813a26953dbdc71c31dacaef9a838c4e29f5712"},
+]
+
 [[package]]
 name = "pillow"
 version = "10.3.0"
@@ -2169,6 +2273,22 @@ tests = ["check-manifest", "coverage", "defusedxml", "markdown2", "olefile", "pa
 typing = ["typing-extensions"]
 xmp = ["defusedxml"]
 
+[[package]]
+name = "platformdirs"
+version = "4.2.1"
+description = "A small Python package for determining appropriate platform-specific dirs, e.g. a `user data dir`."
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "platformdirs-4.2.1-py3-none-any.whl", hash = "sha256:17d5a1161b3fd67b390023cb2d3b026bbd40abde6fdb052dfbd3a29c3ba22ee1"},
+    {file = "platformdirs-4.2.1.tar.gz", hash = "sha256:031cd18d4ec63ec53e82dceaac0417d218a6863f7745dfcc9efe7793b7039bdf"},
+]
+
+[package.extras]
+docs = ["furo (>=2023.9.10)", "proselint (>=0.13)", "sphinx (>=7.2.6)", "sphinx-autodoc-typehints (>=1.25.2)"]
+test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.4.3)", "pytest-cov (>=4.1)", "pytest-mock (>=3.12)"]
+type = ["mypy (>=1.8)"]
+
 [[package]]
 name = "posthog"
 version = "3.5.0"
@@ -3603,4 +3723,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.8.1"
-content-hash = "151d30d44bd27f592b936ec1fd14fff94dcc51746a1440f6c1f36329c75ed744"
+content-hash = "7f3728a66e787b751b7d4564bf9e83cd6fec93b87cedebd5001ecd2a354b7490"

pyproject.toml

@@ -25,8 +25,11 @@ llama-index-embeddings-openai = "~0.1.9"
 llama-index-vector-stores-chroma = "~0.1.7"
 llama-index-cli = "~0.1.12"
 llama-index-readers-file = "~0.1.19"
+pydantic = "~2.7.1"
 
 [tool.poetry.dev-dependencies]
+black = "24.4.2"
+mypy = "1.10.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]