refactor: Separate app in modules

app.py

@@ -3,67 +3,47 @@
 import textwrap
 import os
 import tiktoken
-import chromadb
-from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
 import fitz
 import streamlit as st
 import openai
 from dotenv import load_dotenv
-from openai import OpenAI
+from src.chroma_client import ChromaDB
+import src.gui_messages as gm
+import src.settings as settings
 
 load_dotenv()
 
-chroma_client = chromadb.PersistentClient(path="tmp/chroma")
-chroma_client.heartbeat()
-
-
-def api_message(api_key):
-    """Inform if the api key is set."""
-    if api_key is None:
-        return st.warning("Add your OpenAI API key")
-
-    return st.success("Your API key is setup ")
-
 
 def set_api_key():
     """Set the OpenAI API key."""
     openai.api_key = st.session_state.api_key
-    st.session_state.api_message = api_message(openai.api_key)
+    st.session_state.api_message = gm.api_message(openai.api_key)
 
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 if "api_message" not in st.session_state:
-    st.session_state.api_message = api_message(openai.api_key)
+    st.session_state.api_message = gm.api_message(openai.api_key)
 
-if os.getenv("OPENAI_API_KEY") is None:
-    message = st.session_state.api_message
+# Sidebar
+with st.sidebar:
+    st.write("## OpenAI API key")
     openai.api_key = st.text_input(
-        "Enter your OpenAI API key",
+        "Enter OpenAI API key",
         value="",
         type="password",
         key="api_key",
         placeholder="Enter your OpenAI API key",
         on_change=set_api_key,
         label_visibility="collapsed",
     )
-    st.write("You can find your API key at https://beta.openai.com/account/api-keys")
-    client = OpenAI(api_key=openai.api_key)
-    embedding_function = OpenAIEmbeddingFunction(
-        api_key=openai.api_key, model_name="text-embedding-ada-002"
-    )
-    collection = chroma_client.get_or_create_collection(
-        name="pdf-explainer", embedding_function=embedding_function
-    )
-else:
-    client = OpenAI()
-    embedding_function = OpenAIEmbeddingFunction(
-        api_key=openai.api_key, model_name="text-embedding-ada-002"
-    )
-    collection = chroma_client.get_or_create_collection(
-        name="pdf-explainer", embedding_function=embedding_function
+    st.write(
+        "You can find your API key at https://platform.openai.com/account/api-keys"
     )
 
+# Build settings
+chroma_client = ChromaDB().client
+client, collection = settings.build(chroma_client, api_key=openai.api_key)
 
 # Query ChromaDb
 query = st.text_input(
@@ -90,14 +70,10 @@ def set_api_key():
 
 if pdf is not None:
     with fitz.open(stream=pdf.read(), filetype="pdf") as doc:  # open document
-        text = chr(12).join([page.get_text() for page in doc])
-        st.write(text[0:200])
-        if st.button("Add to collection"):
-            collection.add(
-                documents=[text],
-                metadatas=[{"source": pdf.name}],
-                ids=[pdf.name + str(Date.now())],
-            )
+        with st.spinner("Extracting text..."):
+            text = chr(12).join([page.get_text() for page in doc])
+        st.subheader("Text preview")
+        st.write(text[0:300] + "...")
         if st.button("Save chunks"):
             with st.spinner("Saving chunks..."):
                 chunks = textwrap.wrap(text, 3000)

src/chroma_client.py

@@ -0,0 +1,11 @@
+"""A client for ChromaDB."""
+import chromadb
+
+
+class ChromaDB:
+    """A class for creating a client for ChromaDB."""
+
+    def __init__(self, path="tmp/chroma"):
+        """Initialize the client."""
+        self.client = chromadb.PersistentClient(path=path)
+        self.client.heartbeat()

src/collection_creator.py

@@ -0,0 +1,17 @@
+"""Create a Chroma collection."""
+from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
+import streamlit as st
+
+
+def create_collection(client, name="pdf-explainer", api_key=None):
+    """Create a Chroma collection."""
+    try:
+        embedding_function = OpenAIEmbeddingFunction(
+            api_key=api_key, model_name="text-embedding-ada-002"
+        )
+        collection = client.get_or_create_collection(
+            name=name, embedding_function=embedding_function
+        )
+        return collection
+    except AttributeError as e:
+        st.error("An error ocurred while creating the collection." + e.with_traceback)

src/gui_messages.py

@@ -0,0 +1,18 @@
+"""Streamlit GUI messages."""
+import streamlit as st
+
+
+def api_message(api_key):
+    """Inform if the api key is set."""
+    if api_key is None:
+        return st.warning("Add your OpenAI API key")
+
+    return st.success("Your API key is setup ")
+
+
+def how_to_create_api_key_message(api_key):
+    """Inform how to create an api key."""
+    if api_key is None:
+        return st.write(
+            "You can find your API key at https://beta.openai.com/account/api-keys"
+        )

src/openai_client.py

@@ -0,0 +1,13 @@
+"""OpenAI client creator."""
+import os
+from openai import OpenAI
+
+
+def create_client(api_key=None):
+    """Create an OpenAI client."""
+    if os.getenv("OPENAI_API_KEY"):
+        api_key = os.getenv("OPENAI_API_KEY")
+
+    client = OpenAI(api_key=api_key)
+
+    return client

src/settings.py

@@ -0,0 +1,10 @@
+from src.openai_client import create_client
+from src.collection_creator import create_collection
+
+
+def build(chroma_client, api_key=None):
+    """Build the app."""
+    client = create_client(api_key=api_key)
+    collection = create_collection(chroma_client, api_key=api_key)
+
+    return client, collection