cleanup

hello_wordsmith/hello_wordsmith.py

@@ -1,78 +1,95 @@
-#!/path/to/your/virtualenv/bin/python
+#!/usr/bin/env python
+
+# Standard library imports
 import os
+
+# Third-party imports
+import chromadb
+from llama_index.core import (
+    SimpleDirectoryReader,
+    VectorStoreIndex,
+    StorageContext,
+    PromptTemplate,
+)
 from llama_index.core.ingestion import IngestionPipeline
-from llama_index.core.query_pipeline import QueryPipeline
+from llama_index.core.query_pipeline import QueryPipeline, InputComponent
+from llama_index.core.response_synthesizers import TreeSummarize
 from llama_index.core.storage.docstore import SimpleDocumentStore
-from llama_index.core import SimpleDirectoryReader
 from llama_index.cli.rag import RagCLI
 from llama_index.llms.openai import OpenAI
 from llama_index.vector_stores.chroma import ChromaVectorStore
-import chromadb
 
 
-# optional, set any API keys your script may need (perhaps using python-dotenv library instead)
-# os.environ["OPENAI_API_KEY"] = "sk-xxx"
+def setup_environment():
+    """Set up environment variables, ideally load from .env file for production"""
+    # os.environ["OPENAI_API_KEY"] = "sk-xxx"
+    pass
 
-from llama_index.core import VectorStoreIndex, StorageContext
 
-chroma_client = chromadb.EphemeralClient()
-chroma_collection = chroma_client.create_collection("wordsmith")
-vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
-storage_context = StorageContext.from_defaults(vector_store=vector_store)
+def initialize_chroma_db():
+    """Initialize the ChromaDB client and collection"""
+    chroma_client = chromadb.EphemeralClient()
+    chroma_collection = chroma_client.create_collection("wordsmith")
+    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+    return vector_store
 
 
-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
-)
+def setup_document_storage(vector_store):
+    """Set up document storage and load data"""
+    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()
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    index = VectorStoreIndex.from_documents(docs, storage_context=storage_context)
+    return index
 
-# docstore = SimpleDocumentStore()
 
-llm = OpenAI(api_key=os.environ["OPENAI_API_KEY"], model="gpt-4")
+def initialize_llm():
+    """Initialize the Large Language Model"""
+    llm = OpenAI(api_key=os.environ["OPENAI_API_KEY"], model="gpt-4")
+    return llm
 
-custom_ingestion_pipeline = IngestionPipeline(
-    vector_store=vector_store,
-)
 
-from llama_index.core import PromptTemplate
+def configure_query_pipeline(index, llm):
+    """Configure and set up the query pipeline"""
+    prompt_str = "Please generate related movies to {query_str}"
+    prompt_tmpl = PromptTemplate(prompt_str)
+    query_pipeline = QueryPipeline(verbose=True)
 
+    retriever = index.as_retriever(similarity_top_k=5)
+    summarizer = TreeSummarize(llm=llm, streaming=True)
 
-prompt_str = "Please generate related movies to {query_str}"
-prompt_tmpl = PromptTemplate(prompt_str)
-query_pipeline = QueryPipeline(verbose=True)
+    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")
 
-from llama_index.core.response_synthesizers import TreeSummarize
-from llama_index.core.query_pipeline import InputComponent
-
-# construct vector store and customize storage context
-
-retriever = index.as_retriever(similarity_top_k=5)
-summarizer = TreeSummarize(llm=llm)
-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
 
-# you can optionally specify your own custom readers to support additional file types.
-# file_extractor = {".html": ...}
 
-rag_cli_instance = RagCLI(
-    ingestion_pipeline=custom_ingestion_pipeline,
-    llm=llm,
-    query_pipeline=query_pipeline
-)
+def create_rag_cli(ingestion_pipeline, llm, query_pipeline):
+    """Create the RAG CLI instance"""
+    rag_cli_instance = RagCLI(
+        ingestion_pipeline=ingestion_pipeline, llm=llm, query_pipeline=query_pipeline
+    )
+    return rag_cli_instance
 
 
 def main():
+    setup_environment()
+    vector_store = initialize_chroma_db()
+    index = setup_document_storage(vector_store)
+    llm = initialize_llm()
+    query_pipeline = configure_query_pipeline(index, llm)
+    ingestion_pipeline = IngestionPipeline(vector_store=vector_store)
+    rag_cli_instance = create_rag_cli(ingestion_pipeline, llm, query_pipeline)
     rag_cli_instance.cli()
 
 

setup.py

@@ -1,32 +1,28 @@
 from setuptools import setup, find_packages
 
 setup(
-    name='hello-wordsmith',
-    version='0.1.0',
-    description='A simple Python package to interface with llama-index RAG over wordsmith data.',
-    long_description=open('README.md').read(),
-    long_description_content_type='text/markdown',
-    url='https://huggingface.co/datasets/derek-at-work/test/source env/bin/activate',
-    author='Derek Johnston',
-    author_email='derek@wordsmith.ai',
-    license='MIT',
+    name="hello-wordsmith",
+    version="0.1.0",
+    description="A simple Python package to interface with llama-index RAG over wordsmith data.",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="https://huggingface.co/datasets/derek-at-work/test/source env/bin/activate",
+    author="Derek Johnston",
+    author_email="derek@wordsmith.ai",
+    license="MIT",
     packages=find_packages(),
-    install_requires=[
-        "llama-index",
-        "chromadb",
-        "llama-index-vector-stores-chroma"
-    ],
+    install_requires=["llama-index", "chromadb", "llama-index-vector-stores-chroma"],
     classifiers=[
-        'Intended Audience :: End Users',
-        'Topic :: Software Development :: Build Tools',
-        'License :: OSI Approved :: MIT License',
-        'Programming Language :: Python :: 3',
-        'Programming Language :: Python :: 3.7',
+        "Intended Audience :: End Users",
+        "Topic :: Software Development :: Build Tools",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.7",
     ],
     entry_points={
-        'console_scripts': [
-            'hello-wordsmith=hello_wordsmith.hello_wordsmith:main',
+        "console_scripts": [
+            "hello-wordsmith=hello_wordsmith.hello_wordsmith:main",
         ],
     },
-    python_requires='>=3.6',
+    python_requires=">=3.6",
 )